ksrv/net.go

package ksrv

import (
	"net/http"
	"strconv"
)

type responseWriter struct {
	http.ResponseWriter
	StatueCode int
	err        string
}

func (w *responseWriter) WriteHeader(statusCode int) {
	if w.StatueCode != 0 {
		return
	}
	w.StatueCode = statusCode
	w.ResponseWriter.WriteHeader(statusCode)
}

func (w *responseWriter) Write(body []byte) (int, error) {
	if w.StatueCode >= 500 {
		w.err = string(body)
	}
	if w.StatueCode == 0 {
		w.WriteHeader(200)
	}
	return w.ResponseWriter.Write(body)
}

// Response shorthand for set status code and write body
func Response(w http.ResponseWriter, status int, body []byte) (int, error) {
	w.WriteHeader(status)
	return w.Write(body)
}

// Abort shorthand for aborting with error, strings and status code could also be passed
func Abort(w http.ResponseWriter, errs ...interface{}) (int, error) {
	code := 500
	msg := []byte{}
	for _, err := range errs {
		switch v := err.(type) {
		case Error:
			code = v.Code
			msg = v.Json()
		case int:
			if v >= 100 || v < 600 {
				code = v
			}
		case string:
			msg = []byte(v)
		case error:
			msg = []byte(v.Error())
		}
	}
	if len(msg) == 0 {
		msg = []byte(strconv.Itoa(code))
	}

	w.WriteHeader(code)
	return w.Write(msg)
}