增加ContentLenghtFilter

README.md

@@ -9,6 +9,7 @@
   - string
   - xml
   - json
+  - yaml
 
   内置支持认证方式：
   1. Basic Auth
@@ -28,19 +29,23 @@ go get github.com/xfali/restclient
 ### 基础配置
 
 请参照DefaultRestClient的API说明
-```cassandraql
+```
 //设置读写超时
 func SetTimeout(timeout time.Duration)
 ```
-```cassandraql
+```
 //配置初始转换器列表
 func SetConverters(convs []Converter)
 ```
-```cassandraql
+```
 //配置连接池
 func SetRoundTripper(tripper http.RoundTripper)
 ```
-```cassandraql
+```
+// 增加处理filter
+func AddFilter(filters ...Filter)
+```
+```
 //配置request创建器
 func SetRequestCreator(f RequestCreator)
 ```
@@ -50,7 +55,7 @@ func SetRequestCreator(f RequestCreator)
 
 ## 使用
 
-```cassandraql
+```
 //使用默认配置
 client := restclient.New()
 str := ""

buffer/buffer.go

@@ -8,7 +8,6 @@ package buffer
 import (
 	"bytes"
 	"io"
-	"io/ioutil"
 	"sync"
 )
 
@@ -72,38 +71,70 @@ func (p *defaultPool) Put(buf *bytes.Buffer) {
 	p.pool.Put(buf)
 }
 
-func NewReadCloser(d []byte) io.ReadCloser {
+type ContentLength interface {
+	ContentLength() int64
+}
+
+type NopReadCloser struct {
+	r   io.Reader
+	len int
+}
+
+func (rc *NopReadCloser) Read(p []byte) (n int, err error) {
+	return rc.r.Read(p)
+}
+
+func (rc *NopReadCloser) Close() error {
+	return nil
+}
+
+func (rc *NopReadCloser) ContentLength() int64 {
+	return int64(rc.len)
+}
+
+func NewReadCloser(d []byte) *NopReadCloser {
 	if d == nil {
 		return nil
 	}
-	return ioutil.NopCloser(bytes.NewReader(d))
+	return &NopReadCloser{
+		r:   bytes.NewReader(d),
+		len: len(d),
+	}
 }
 
-type readWriteCloser struct {
+type ReadWriteCloser struct {
 	pool Pool
 	buf  *bytes.Buffer
+	once sync.Once
 }
 
-func (rc *readWriteCloser) Bytes() []byte {
+func (rc *ReadWriteCloser) Bytes() []byte {
 	return rc.buf.Bytes()
 }
 
-func (rc *readWriteCloser) Read(p []byte) (n int, err error) {
+func (rc *ReadWriteCloser) Read(p []byte) (n int, err error) {
 	return rc.buf.Read(p)
 }
 
-func (rc *readWriteCloser) Write(p []byte) (n int, err error) {
+func (rc *ReadWriteCloser) Write(p []byte) (n int, err error) {
 	return rc.buf.Write(p)
 }
 
-func (rc *readWriteCloser) Close() error {
-	rc.pool.Put(rc.buf)
+func (rc *ReadWriteCloser) Close() error {
+	// just return once
+	rc.once.Do(func() {
+		rc.pool.Put(rc.buf)
+	})
 	return nil
 }
 
-func NewReadWriteCloser(pool Pool) *readWriteCloser {
+func (rc *ReadWriteCloser) ContentLength() int64 {
+	return int64(rc.buf.Len())
+}
+
+func NewReadWriteCloser(pool Pool) *ReadWriteCloser {
 	buf := pool.Get()
-	return &readWriteCloser{
+	return &ReadWriteCloser{
 		pool: pool,
 		buf:  buf,
 	}

wrap_client.go

@@ -17,6 +17,7 @@ import (
 	"net/http"
 	"net/url"
 	"reflect"
+	"strconv"
 	"time"
 )
 
@@ -163,6 +164,8 @@ func (auth *DigestAuth) Filter(request *http.Request, fc FilterChain) (*http.Res
 			return nil, err
 		}
 		reqData = buf.Bytes()
+		// close old request body
+		request.Body.Close()
 		request.Body = buffer.NewReadCloser(reqData)
 	}
 
@@ -214,6 +217,24 @@ func findWWWAuth(header http.Header) string {
 	return ""
 }
 
+func ContentLengthFilter(request *http.Request, fc FilterChain) (*http.Response, error) {
+	if request.ContentLength > 0 {
+		return fc.Filter(request)
+	}
+	lengthStr := request.Header.Get("Content-Length")
+	if lengthStr != "" {
+		l, err := strconv.ParseInt(lengthStr, 10, 64)
+		if err == nil {
+			request.ContentLength = l
+		}
+		return fc.Filter(request)
+	}
+	if cl, ok := request.Body.(buffer.ContentLength); ok {
+		request.ContentLength = cl.ContentLength()
+	}
+	return fc.Filter(request)
+}
+
 type LogFunc func(format string, args ...interface{})
 type Log struct {
 	Log  xlog.Logger
@@ -243,6 +264,8 @@ func (log *Log) Filter(request *http.Request, fc FilterChain) (*http.Response, e
 			return nil, err
 		}
 		reqData = buf.Bytes()
+		// close old request body
+		request.Body.Close()
 		request.Body = buffer.NewReadCloser(reqData)
 	}