sillygod/cdp-cache

package httpcache

import (
    "errors"
    "io"
    "net/http"
    "strings"

    "github.com/sillygod/cdp-cache/backends"
)

// copyHeaders copy the header from one to another
func copyHeaders(from http.Header, to http.Header) {
    for k, values := range from {
        for _, v := range values {
            to.Set(k, v)
        }
    }
}

// Response encapsulates the entry
type Response struct {
    Code       int
    HeaderMap  http.Header
    body       backends.Backend
    snapHeader http.Header

    wroteHeader        bool
    bodyComplete       bool
    IsFirstByteWritten bool

    bodyChan         chan struct{} // indicates whether the backend is set or not.
    bodyCompleteChan chan struct{}
    closedChan       chan struct{}
    headersChan      chan struct{}
}

// NewResponse returns an initialized Response.
func NewResponse() *Response {
    r := &Response{
        Code:             200,
        HeaderMap:        http.Header{},
        body:             nil,
        bodyChan:         make(chan struct{}, 1),
        closedChan:       make(chan struct{}, 1),
        headersChan:      make(chan struct{}, 1),
        bodyCompleteChan: make(chan struct{}, 1),
    }
    return r
}

// Header return the header from the upstream response
func (r *Response) Header() http.Header {
    return r.HeaderMap
}

func (r *Response) writeHeader(b []byte, str string) {
    if r.wroteHeader {
        return
    }

    if len(str) > 512 {
        str = str[:512]
    }

    h := r.Header()
    _, hasType := h["Content-Type"]
    hasTE := h.Get("Transfer-Encoding") != ""

    if !hasType && !hasTE {
        if b == nil {
            b = []byte(str)
        }

        if len(b) > 512 {
            b = b[:512]
        }

        h.Set("Content-Type", http.DetectContentType(b))
    }

    r.WriteHeader(r.Code)
}

// Write writes the upstream's content in the backend's storage
// this will wait the body(backend) is set
func (r *Response) Write(buf []byte) (int, error) {

    if !r.wroteHeader {
        r.writeHeader(buf, "")
    }

    if r.body == nil {
        <-r.bodyChan
    }

    if r.body != nil {
        if !r.IsFirstByteWritten {
            r.IsFirstByteWritten = true
        }
        return r.body.Write(buf)
    }

    return 0, errors.New("No storage provided")
}

// WaitClose waits the response to be closed.
func (r *Response) WaitClose() {
    <-r.closedChan
}

// GetReader gets the reader from the setted backend
func (r *Response) GetReader() (io.ReadCloser, error) {
    if r.bodyComplete == false {
        <-r.bodyCompleteChan
    }
    return r.body.GetReader()
}

// SetBody sets the backend to body for the further write usage
func (r *Response) SetBody(body backends.Backend) {
    r.body = body
    r.bodyChan <- struct{}{}
}

// WaitHeaders waits the header to be written
func (r *Response) WaitHeaders() {
    <-r.headersChan
}

// WriteHeader keeps the upstream response header
func (r *Response) WriteHeader(code int) {
    if r.wroteHeader {
        return
    }

    r.Code = code
    r.wroteHeader = true

    r.snapHeader = http.Header{}
    copyHeaders(r.Header(), r.snapHeader)
    r.snapHeader.Del("server")

    r.headersChan <- struct{}{}
}

func shouldUseCache(req *http.Request, config *Config) bool {

    matchMethod := false
    for _, method := range config.MatchMethods {
        if method == req.Method {
            matchMethod = true
            break
        }
    }
    if !matchMethod {
        return false
    }

    // Range requests still not supported
    // It may happen that the previous request for this url has a successful response
    // but for another Range. So a special handling is needed
    if req.Header.Get("range") != "" {
        return false
    }

    if isWebSocket(req.Header) {
        return false
    }

    return true

}

func isWebSocket(h http.Header) bool {
    if h == nil {
        return false
    }

    // Get gets the *first* value associated with the given key.
    if strings.ToLower(h.Get("Upgrade")) != "websocket" {
        return false
    }

    // To access multiple values of a key, access the map directly.
    for _, value := range h.Values("Connection") {
        if strings.ToLower(value) == "upgrade" {
            return true
        }
    }

    return false
}

// Flush flushes the backend's storage (currently, only file storage need to call this)
func (r *Response) Flush() {
    if !r.wroteHeader {
        r.WriteHeader(200)
    }

    if r.body == nil {
        return
    }

    r.body.Flush()
}

// Close indicate the data is completely written to the body
// so that we can close it.
func (r *Response) Close() error {
    if r.body == nil {
        <-r.bodyChan
    }

    r.body.Close()

    r.bodyComplete = true
    r.bodyCompleteChan <- struct{}{}
    r.closedChan <- struct{}{}

    return nil
}

// Clean performs purge the cache
func (r *Response) Clean() error {

    if r.body == nil {
        return nil
    }

    return r.body.Clean()
}