141 lines
2.9 KiB
Go
141 lines
2.9 KiB
Go
package tango
|
|
|
|
import (
|
|
"bufio"
|
|
"compress/flate"
|
|
"compress/gzip"
|
|
"fmt"
|
|
"io"
|
|
"net"
|
|
"net/http"
|
|
"path"
|
|
"strings"
|
|
)
|
|
|
|
const (
|
|
HeaderAcceptEncoding = "Accept-Encoding"
|
|
HeaderContentEncoding = "Content-Encoding"
|
|
HeaderContentLength = "Content-Length"
|
|
HeaderContentType = "Content-Type"
|
|
HeaderVary = "Vary"
|
|
)
|
|
|
|
type Compresser interface {
|
|
CompressType() string
|
|
}
|
|
|
|
type GZip struct {}
|
|
func (GZip) CompressType() string {
|
|
return "gzip"
|
|
}
|
|
|
|
type Deflate struct {}
|
|
func (Deflate) CompressType() string {
|
|
return "deflate"
|
|
}
|
|
|
|
type Compress struct {}
|
|
func (Compress) CompressType() string {
|
|
return "auto"
|
|
}
|
|
|
|
func Compresses(exts []string) HandlerFunc{
|
|
extsmap := make(map[string]bool)
|
|
for _, ext := range exts {
|
|
extsmap[strings.ToLower(ext)] = true
|
|
}
|
|
|
|
return func(ctx *Context) {
|
|
ae := ctx.Req().Header.Get("Accept-Encoding")
|
|
if ae == "" {
|
|
ctx.Next()
|
|
return
|
|
}
|
|
|
|
if len(extsmap) > 0 {
|
|
ext := strings.ToLower(path.Ext(ctx.Req().URL.Path))
|
|
if _, ok := extsmap[ext]; ok {
|
|
compress(ctx, "auto")
|
|
return
|
|
}
|
|
}
|
|
|
|
if action := ctx.Action(); action != nil {
|
|
if c, ok := action.(Compresser); ok {
|
|
compress(ctx, c.CompressType())
|
|
return
|
|
}
|
|
}
|
|
|
|
// if blank, then no compress
|
|
ctx.Next()
|
|
}
|
|
}
|
|
|
|
func compress(ctx *Context, compressType string) {
|
|
ae := ctx.Req().Header.Get("Accept-Encoding")
|
|
acceptCompress := strings.SplitN(ae, ",", -1)
|
|
var writer io.Writer
|
|
var val string
|
|
|
|
for _, val = range acceptCompress {
|
|
val = strings.TrimSpace(val)
|
|
if compressType == "auto" || val == compressType {
|
|
if val == "gzip" {
|
|
ctx.Header().Set("Content-Encoding", "gzip")
|
|
writer = gzip.NewWriter(ctx.ResponseWriter)
|
|
break
|
|
} else if val == "deflate" {
|
|
ctx.Header().Set("Content-Encoding", "deflate")
|
|
writer, _ = flate.NewWriter(ctx.ResponseWriter, flate.BestSpeed)
|
|
break
|
|
}
|
|
}
|
|
}
|
|
|
|
// not supported compress method, then ignore
|
|
if writer == nil {
|
|
ctx.Next()
|
|
return
|
|
}
|
|
|
|
// for cache server
|
|
ctx.Header().Add(HeaderVary, "Accept-Encoding")
|
|
|
|
gzw := &compressWriter{writer, ctx.ResponseWriter}
|
|
ctx.ResponseWriter = gzw
|
|
|
|
ctx.Next()
|
|
|
|
// delete content length after we know we have been written to
|
|
gzw.Header().Del(HeaderContentLength)
|
|
ctx.ResponseWriter = gzw.ResponseWriter
|
|
|
|
switch writer.(type) {
|
|
case *gzip.Writer:
|
|
writer.(*gzip.Writer).Close()
|
|
case *flate.Writer:
|
|
writer.(*flate.Writer).Close()
|
|
}
|
|
}
|
|
|
|
type compressWriter struct {
|
|
w io.Writer
|
|
ResponseWriter
|
|
}
|
|
|
|
func (grw *compressWriter) Write(p []byte) (int, error) {
|
|
if len(grw.Header().Get(HeaderContentType)) == 0 {
|
|
grw.Header().Set(HeaderContentType, http.DetectContentType(p))
|
|
}
|
|
return grw.w.Write(p)
|
|
}
|
|
|
|
func (grw *compressWriter) Hijack() (net.Conn, *bufio.ReadWriter, error) {
|
|
hijacker, ok := grw.ResponseWriter.(http.Hijacker)
|
|
if !ok {
|
|
return nil, nil, fmt.Errorf("the ResponseWriter doesn't support the Hijacker interface")
|
|
}
|
|
return hijacker.Hijack()
|
|
}
|