apinto/node/http-context/header.go

package http_context

import (
	"bytes"
	"fmt"
	"net/http"
	"strings"
	"sync"

	http_service "github.com/eolinker/eosc/eocontext/http-context"

	"github.com/valyala/fasthttp"
)

var _ http_service.IHeaderWriter = (*RequestHeader)(nil)

type RequestHeader struct {
	header *fasthttp.RequestHeader
	tmp    http.Header
}

func (h *RequestHeader) RawHeader() string {
	return h.header.String()
}

func (h *RequestHeader) reset(header *fasthttp.RequestHeader) {
	h.header = header
	h.tmp = nil
}

func (h *RequestHeader) initHeader() {
	if h.tmp == nil {
		h.tmp = make(http.Header)
		h.header.VisitAll(func(key, value []byte) {
			bytes.SplitN(value, []byte(":"), 2)
			h.tmp[string(key)] = []string{string(value)}
		})
	}
}

func (h *RequestHeader) Host() string {
	return string(h.header.Host())
}

func (h *RequestHeader) GetHeader(name string) string {
	return h.Headers().Get(name)
}

func (h *RequestHeader) Headers() http.Header {
	h.initHeader()
	return h.tmp
}

func (h *RequestHeader) SetHeader(key, value string) {
	if h.tmp != nil {
		h.tmp.Set(key, value)
	}
	h.header.Set(key, value)
}

func (h *RequestHeader) AddHeader(key, value string) {
	if h.tmp != nil {
		h.tmp.Add(key, value)
	}
	h.header.Add(key, value)
}

func (h *RequestHeader) DelHeader(key string) {
	if h.tmp != nil {
		h.tmp.Del(key)
	}
	h.header.Del(key)
}

func (h *RequestHeader) SetHost(host string) {
	if h.tmp != nil {
		h.tmp.Set("Host", host)
	}
	h.header.SetHost(host)
}

type headerActionHandleFunc func(target *ResponseHeader, key string, value ...string)

var (
	headerActionAdd = func(target *ResponseHeader, key string, value ...string) {
		target.cache.Add(key, value[0])
		target.header.Add(key, value[0])

	}
	headerActionSet = func(target *ResponseHeader, key string, value ...string) {
		target.cache.Set(key, value[0])
		target.header.Set(key, value[0])
	}
	headerActionDel = func(target *ResponseHeader, key string, value ...string) {
		target.cache.Del(key)
		target.header.Del(key)
	}
)

type headerAction struct {
	Action headerActionHandleFunc
	Key    string
	Value  string
}
type ResponseHeader struct {
	header     *fasthttp.ResponseHeader
	locker     sync.RWMutex
	cache      http.Header
	actions    []*headerAction
	afterProxy bool
}

func (r *ResponseHeader) reset(header *fasthttp.ResponseHeader) {

	r.header = header
	r.cache = http.Header{}
	r.actions = nil
	r.afterProxy = false
	//r.refresh()
}
func (r *ResponseHeader) refresh() {
	tmp := make(http.Header)
	hs := strings.Split(r.header.String(), "\r\n")
	for i, t := range hs {
		if i == 0 {
			continue
		}
		if strings.TrimSpace(t) == "" {
			continue
		}
		vs := strings.SplitN(t, ":", 2)
		if len(vs) < 2 {
			if vs[0] == "" {
				continue
			}
			tmp[vs[0]] = []string{""}
			continue
		}
		tmp[vs[0]] = []string{strings.TrimSpace(vs[1])}
	}
	r.cache = tmp
	for _, ac := range r.actions {
		ac.Action(r, ac.Key, ac.Value)
	}
	r.afterProxy = true
	r.actions = nil

}
func (r *ResponseHeader) Finish() {
	r.header = nil
	r.cache = nil
	r.actions = nil
}

func (r *ResponseHeader) GetHeader(name string) string {
	return r.Headers().Get(name)
}

func (r *ResponseHeader) Headers() http.Header {
	r.locker.RLock()
	defer r.locker.RUnlock()
	return r.cache.Clone()
}

func (r *ResponseHeader) SetHeader(key, value string) {
	r.locker.Lock()
	defer r.locker.Unlock()
	r.cache.Set(key, value)
	r.header.Set(key, value)
	if !r.afterProxy {
		r.actions = append(r.actions, &headerAction{
			Key:    key,
			Value:  value,
			Action: headerActionSet,
		})
	}
}

func (r *ResponseHeader) AddHeader(key, value string) {
	r.locker.Lock()
	defer r.locker.Unlock()
	r.cache.Add(key, value)
	r.header.Add(key, value)
	if !r.afterProxy {
		r.actions = append(r.actions, &headerAction{
			Key:    key,
			Value:  value,
			Action: headerActionAdd,
		})
	}
}

func (r *ResponseHeader) DelHeader(key string) {
	r.locker.Lock()
	defer r.locker.Unlock()
	r.cache.Del(key)
	r.header.Del(key)
	if !r.afterProxy {
		r.actions = append(r.actions, &headerAction{
			Key:    key,
			Action: headerActionDel,
		})
	}

}

func (h *RequestHeader) GetCookie(key string) string {
	return string(h.header.Cookie(key))
}

func (h *RequestHeader) SetCookie(key, value string, maxAge int) {
	h.header.SetCookie(key, value)

	if maxAge > 0 {
		h.header.SetCookie(key, fmt.Sprintf("%s; Max-Age=%d", value, maxAge))
	}
}