Skip to content
Permalink
Branch: master
Find file Copy path
Find file Copy path
Fetching contributors…
Cannot retrieve contributors at this time
115 lines (98 sloc) 2.28 KB
package torproxy
import (
"bytes"
"compress/gzip"
"io"
"net/http"
"sync"
"golang.org/x/net/proxy"
)
type TorResponse struct {
headers http.Header
body []byte
bodyReader bytes.Buffer
bodyWriter bytes.Buffer
status int
request *http.Request
dialer proxy.Dialer
}
var bufferPool = sync.Pool{New: createBuffer}
func createBuffer() interface{} {
return make([]byte, 0, 32*1024)
}
// Header returns response headers
func (t *TorResponse) Header() http.Header {
return t.headers
}
func (t *TorResponse) Write(body []byte) (int, error) {
reader := bytes.NewReader(body)
pooledIoCopy(&t.bodyReader, reader)
t.body = body
return len(body), nil
}
// Body returns response's body. This method should only get called after WriteBody()
func (t *TorResponse) Body() []byte {
return t.bodyWriter.Bytes()
}
// WriteHeader Writes the given status code to response
func (t *TorResponse) WriteHeader(status int) {
t.status = status
}
func (t *TorResponse) ReplaceBody(scheme, to, host string) error {
replacedBody := bytes.Replace(t.bodyWriter.Bytes(), []byte(scheme+"://"+to), []byte(scheme+"://"+host), -1)
t.bodyWriter.Reset()
if _, err := t.bodyWriter.Write(replacedBody); err != nil {
return err
}
return nil
}
func (t *TorResponse) WriteBody() error {
switch t.Header().Get("Content-Encoding") {
case "gzip":
reader, err := gzip.NewReader(&t.bodyReader)
if err != nil {
return err
}
defer reader.Close()
_, err = io.Copy(&t.bodyWriter, reader)
if err != nil {
return err
}
t.Header().Del("Content-Encoding")
default:
_, err := io.Copy(&t.bodyWriter, &t.bodyReader)
if err != nil {
return err
}
}
return nil
}
var skipHeaders = map[string]struct{}{
"Content-Type": {},
"Content-Disposition": {},
"Accept-Ranges": {},
"Set-Cookie": {},
"Cache-Control": {},
"Expires": {},
}
func copyHeader(dst, src http.Header) {
for k, vv := range src {
if _, ok := dst[k]; ok {
if _, shouldSkip := skipHeaders[k]; shouldSkip {
continue
}
if k != "Server" {
dst.Del(k)
}
}
for _, v := range vv {
dst.Add(k, v)
}
}
}
func pooledIoCopy(dst io.Writer, src io.Reader) {
buf := bufferPool.Get().([]byte)
defer bufferPool.Put(buf)
bufCap := cap(buf)
io.CopyBuffer(dst, src, buf[0:bufCap:bufCap])
}
You can’t perform that action at this time.