websocket/compression.go

141 lines
2.8 KiB
Go
Raw Normal View History

2016-12-16 01:53:35 +03:00
// Copyright 2017 The Gorilla WebSocket Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.
package websocket
import (
"compress/flate"
"errors"
"io"
"strings"
2016-12-16 01:53:35 +03:00
"sync"
)
var (
flateWriterPool = sync.Pool{}
2016-12-18 02:33:06 +03:00
flateReaderPool = sync.Pool{}
)
2016-12-18 02:33:06 +03:00
func decompressNoContextTakeover(r io.Reader) io.ReadCloser {
const tail =
// Add four bytes as specified in RFC
"\x00\x00\xff\xff" +
// Add final block to squelch unexpected EOF error from flate reader.
"\x01\x00\x00\xff\xff"
2016-12-18 02:33:06 +03:00
i := flateReaderPool.Get()
if i == nil {
i = flate.NewReader(nil)
}
i.(flate.Resetter).Reset(io.MultiReader(r, strings.NewReader(tail)), nil)
return &flateReadWrapper{i.(io.ReadCloser)}
}
func compressNoContextTakeover(w io.WriteCloser) (io.WriteCloser, error) {
tw := &truncWriter{w: w}
2016-12-16 01:53:35 +03:00
i := flateWriterPool.Get()
var fw *flate.Writer
var err error
if i == nil {
fw, err = flate.NewWriter(tw, 3)
} else {
fw = i.(*flate.Writer)
fw.Reset(tw)
}
2016-12-18 02:33:06 +03:00
return &flateWriteWrapper{fw: fw, tw: tw}, err
}
// truncWriter is an io.Writer that writes all but the last four bytes of the
// stream to another io.Writer.
type truncWriter struct {
w io.WriteCloser
n int
p [4]byte
}
func (w *truncWriter) Write(p []byte) (int, error) {
n := 0
// fill buffer first for simplicity.
if w.n < len(w.p) {
n = copy(w.p[w.n:], p)
p = p[n:]
w.n += n
if len(p) == 0 {
return n, nil
}
}
m := len(p)
if m > len(w.p) {
m = len(w.p)
}
if nn, err := w.w.Write(w.p[:m]); err != nil {
return n + nn, err
}
copy(w.p[:], w.p[m:])
copy(w.p[len(w.p)-m:], p[len(p)-m:])
nn, err := w.w.Write(p[:len(p)-m])
return n + nn, err
}
2016-12-18 02:33:06 +03:00
type flateWriteWrapper struct {
fw *flate.Writer
tw *truncWriter
}
2016-12-18 02:33:06 +03:00
func (w *flateWriteWrapper) Write(p []byte) (int, error) {
2016-12-16 01:53:35 +03:00
if w.fw == nil {
return 0, errWriteClosed
}
return w.fw.Write(p)
}
2016-12-18 02:33:06 +03:00
func (w *flateWriteWrapper) Close() error {
2016-12-16 01:53:35 +03:00
if w.fw == nil {
return errWriteClosed
}
err1 := w.fw.Flush()
2016-12-16 01:53:35 +03:00
flateWriterPool.Put(w.fw)
w.fw = nil
if w.tw.p != [4]byte{0, 0, 0xff, 0xff} {
return errors.New("websocket: internal error, unexpected bytes at end of flate stream")
}
err2 := w.tw.w.Close()
if err1 != nil {
return err1
}
return err2
}
2016-12-18 02:33:06 +03:00
type flateReadWrapper struct {
fr io.ReadCloser
}
func (r *flateReadWrapper) Read(p []byte) (int, error) {
if r.fr == nil {
return 0, io.ErrClosedPipe
}
n, err := r.fr.Read(p)
if err == io.EOF {
// Preemptively place the reader back in the pool. This helps with
// scenarios where the application does not call NextReader() soon after
// this final read.
r.Close()
}
return n, err
}
func (r *flateReadWrapper) Close() error {
if r.fr == nil {
return io.ErrClosedPipe
}
err := r.fr.Close()
flateReaderPool.Put(r.fr)
r.fr = nil
return err
}