Browse Source

update http/tcp/aes

ls 2 years ago
parent
commit
5f22a4200c
13 changed files with 212 additions and 54 deletions
  1. 7 1
      crypto/aes.go
  2. 4 0
      db/init.go
  3. 25 0
      pool/buffer.go
  4. 26 0
      pool/buffer_test.go
  5. 1 0
      tcp/error.go
  6. 1 0
      tcp/optimize.go
  7. 1 0
      tcp/read.go
  8. 1 0
      tcp/server.go
  9. 16 0
      tcp/write.go
  10. 88 52
      xhttp/client.go
  11. 17 0
      xhttp/option.go
  12. 24 0
      xhttp/response.go
  13. 1 1
      xhttp/serve.go

+ 7 - 1
crypto/aes.go

@@ -4,6 +4,7 @@ import (
 	"bytes"
 	"bytes"
 	"crypto/aes"
 	"crypto/aes"
 	"crypto/cipher"
 	"crypto/cipher"
+	"errors"
 )
 )
 
 
 // AES aes
 // AES aes
@@ -52,9 +53,14 @@ func (c *AES) Decrypt(bs []byte) (ds []byte, err error) {
 	if err != nil {
 	if err != nil {
 		return nil, err
 		return nil, err
 	}
 	}
+	n := len(bs)
 	blockSize := block.BlockSize()
 	blockSize := block.BlockSize()
+	if n%blockSize != 0 {
+		return nil, errors.New(`crypto/cipher: input not full blocks`)
+	}
+
 	blockMode := cipher.NewCBCDecrypter(block, c.iv[:blockSize])
 	blockMode := cipher.NewCBCDecrypter(block, c.iv[:blockSize])
-	ds = make([]byte, len(bs))
+	ds = make([]byte, n)
 	blockMode.CryptBlocks(ds, bs)
 	blockMode.CryptBlocks(ds, bs)
 	ds = pkcs7UnPadding(ds)
 	ds = pkcs7UnPadding(ds)
 	return
 	return

+ 4 - 0
db/init.go

@@ -0,0 +1,4 @@
+package db
+
+func init() {
+}

+ 25 - 0
pool/buffer.go

@@ -0,0 +1,25 @@
+package pool
+
+import (
+	"bytes"
+	"sync"
+)
+
+var bufferPool = sync.Pool{
+	New: func() interface{} {
+		return &bytes.Buffer{}
+	},
+}
+
+// NewBuffer get Buffer from poll
+func NewBuffer() *bytes.Buffer {
+	b := bufferPool.Get().(*bytes.Buffer)
+	b.Reset()
+	return b
+}
+
+// PutBuffer add *bytes.Buffer to the pool.
+func PutBuffer(b *bytes.Buffer) {
+	b.Reset()
+	bufferPool.Put(b)
+}

+ 26 - 0
pool/buffer_test.go

@@ -0,0 +1,26 @@
+package pool_test
+
+import (
+	"bytes"
+	"testing"
+
+	"git.chuangxin1.com/myth/sacred/pool"
+)
+
+var data = make([]byte, 10000)
+
+func BenchmarkBufferWithPool(b *testing.B) {
+	for n := 0; n < b.N; n++ {
+		buf := pool.NewBuffer()
+		buf.Write(data)
+		buf.Reset()
+		pool.PutBuffer(buf)
+	}
+}
+
+func BenchmarkBuffer(b *testing.B) {
+	for n := 0; n < b.N; n++ {
+		var buf bytes.Buffer
+		buf.Write(data)
+	}
+}

+ 1 - 0
tcp/error.go

@@ -0,0 +1 @@
+package tcp

+ 1 - 0
tcp/optimize.go

@@ -0,0 +1 @@
+package tcp

+ 1 - 0
tcp/read.go

@@ -0,0 +1 @@
+package tcp

+ 1 - 0
tcp/server.go

@@ -0,0 +1 @@
+package tcp

+ 16 - 0
tcp/write.go

@@ -0,0 +1,16 @@
+package tcp
+
+import (
+	"net"
+	"time"
+)
+
+// Write write to socket with timeout
+func Write(c *net.TCPConn, t time.Duration, b []byte) (n int, err error) {
+	err = c.SetWriteDeadline(time.Now().Add(t))
+	if err != nil {
+		return
+	}
+	n, err = c.Write(b)
+	return
+}

+ 88 - 52
xhttp/client.go

@@ -11,12 +11,19 @@ import (
 )
 )
 
 
 // readBody read response
 // readBody read response
-func readBody(res *http.Response) (msg Message, err error) {
+func readBody(req *http.Request, res *http.Response) (msg Message, err error) {
 	var (
 	var (
 		body   []byte
 		body   []byte
 		reader io.Reader
 		reader io.Reader
 	)
 	)
 
 
+	msg.StatusCode = res.StatusCode
+	msg.Header = res.Header
+
+	defer res.Body.Close()
+	msg.res = res
+	msg.req = req
+
 	encoding := res.Header.Get("Content-Encoding")
 	encoding := res.Header.Get("Content-Encoding")
 	switch encoding {
 	switch encoding {
 	case "gzip":
 	case "gzip":
@@ -31,14 +38,12 @@ func readBody(res *http.Response) (msg Message, err error) {
 		return
 		return
 	}
 	}
 
 
-	msg.StatusCode = res.StatusCode
-	msg.Header = res.Header
 	msg.Body = body
 	msg.Body = body
 	return
 	return
 }
 }
 
 
 // newRequest new request
 // newRequest new request
-func newRequest(method, uri string, body io.Reader, opt RequestOption) (res *http.Response, err error) {
+func newRequest(method, uri string, body io.Reader, opt RequestOption) (req *http.Request, res *http.Response, err error) {
 	t := &http.Transport{
 	t := &http.Transport{
 		Dial: func(netw, addr string) (net.Conn, error) {
 		Dial: func(netw, addr string) (net.Conn, error) {
 			conn, err := net.DialTimeout(netw, addr, opt.RequestTimeOut)
 			conn, err := net.DialTimeout(netw, addr, opt.RequestTimeOut)
@@ -62,7 +67,7 @@ func newRequest(method, uri string, body io.Reader, opt RequestOption) (res *htt
 		t.TLSClientConfig = &tls.Config{InsecureSkipVerify: true}
 		t.TLSClientConfig = &tls.Config{InsecureSkipVerify: true}
 	}
 	}
 	c := &http.Client{Transport: t}
 	c := &http.Client{Transport: t}
-	var req *http.Request
+	//var req *http.Request
 	req, err = http.NewRequest(method, uri, body)
 	req, err = http.NewRequest(method, uri, body)
 	if err != nil {
 	if err != nil {
 		return
 		return
@@ -72,43 +77,51 @@ func newRequest(method, uri string, body io.Reader, opt RequestOption) (res *htt
 		req.Header.Add(k, v)
 		req.Header.Add(k, v)
 	}
 	}
 
 
+	if opt.Cookies != nil {
+		for _, c := range opt.Cookies {
+			req.AddCookie(c)
+		}
+	}
+
+	if len(opt.username) > 0 {
+		req.SetBasicAuth(opt.username, opt.password)
+	}
+
 	res, err = c.Do(req)
 	res, err = c.Do(req)
 	return
 	return
 }
 }
 
 
 // Get HTTP request GET
 // Get HTTP request GET
 func Get(uri string, opt RequestOption) (msg Message, err error) {
 func Get(uri string, opt RequestOption) (msg Message, err error) {
-	var res *http.Response
-	if res, err = newRequest("GET", uri, nil, opt); err != nil {
+	var (
+		res *http.Response
+		req *http.Request
+	)
+	if req, res, err = newRequest("GET", uri, nil, opt); err != nil {
 		return
 		return
 	}
 	}
 
 
-	defer res.Body.Close()
-	msg, err = readBody(res)
+	msg, err = readBody(req, res)
 	return
 	return
 }
 }
 
 
 // Head HTTP request HEAD
 // Head HTTP request HEAD
 func Head(uri string, opt RequestOption) (msg Message, err error) {
 func Head(uri string, opt RequestOption) (msg Message, err error) {
-	var res *http.Response
-	if res, err = newRequest("HEAD", uri, nil, opt); err != nil {
-		return
-	}
-
-	defer res.Body.Close()
-	//msg, err = readBody(res)
+	msg.req, msg.res, err = newRequest("HEAD", uri, nil, opt)
 	return
 	return
 }
 }
 
 
 // GetJSON HTTP request GET, response JSON
 // GetJSON HTTP request GET, response JSON
 func GetJSON(v interface{}, uri string, opt RequestOption) (msg Message, err error) {
 func GetJSON(v interface{}, uri string, opt RequestOption) (msg Message, err error) {
-	var res *http.Response
-	if res, err = newRequest("GET", uri, nil, opt); err != nil {
+	var (
+		res *http.Response
+		req *http.Request
+	)
+	if req, res, err = newRequest("GET", uri, nil, opt); err != nil {
 		return
 		return
 	}
 	}
 
 
-	defer res.Body.Close()
-	msg, err = readBody(res)
+	msg, err = readBody(req, res)
 	if err != nil {
 	if err != nil {
 		return
 		return
 	}
 	}
@@ -118,13 +131,15 @@ func GetJSON(v interface{}, uri string, opt RequestOption) (msg Message, err err
 
 
 // GetXML HTTP request GET, response XML
 // GetXML HTTP request GET, response XML
 func GetXML(v interface{}, uri string, opt RequestOption) (msg Message, err error) {
 func GetXML(v interface{}, uri string, opt RequestOption) (msg Message, err error) {
-	var res *http.Response
-	if res, err = newRequest("GET", uri, nil, opt); err != nil {
+	var (
+		res *http.Response
+		req *http.Request
+	)
+	if req, res, err = newRequest("GET", uri, nil, opt); err != nil {
 		return
 		return
 	}
 	}
 
 
-	defer res.Body.Close()
-	msg, err = readBody(res)
+	msg, err = readBody(req, res)
 	if err != nil {
 	if err != nil {
 		return
 		return
 	}
 	}
@@ -134,77 +149,98 @@ func GetXML(v interface{}, uri string, opt RequestOption) (msg Message, err erro
 
 
 // Post HTTP request POST
 // Post HTTP request POST
 func Post(uri string, body io.Reader, opt RequestOption) (msg Message, err error) {
 func Post(uri string, body io.Reader, opt RequestOption) (msg Message, err error) {
-	var res *http.Response
-	if res, err = newRequest("POST", uri, body, opt); err != nil {
+	var (
+		res *http.Response
+		req *http.Request
+	)
+	if req, res, err = newRequest("POST", uri, body, opt); err != nil {
 		return
 		return
 	}
 	}
-	defer res.Body.Close()
-	msg, err = readBody(res)
+
+	msg, err = readBody(req, res)
 	return
 	return
 }
 }
 
 
 // Put HTTP request PUT
 // Put HTTP request PUT
 func Put(uri string, body io.Reader, opt RequestOption) (msg Message, err error) {
 func Put(uri string, body io.Reader, opt RequestOption) (msg Message, err error) {
-	var res *http.Response
-	if res, err = newRequest("PUT", uri, body, opt); err != nil {
+	var (
+		res *http.Response
+		req *http.Request
+	)
+	if req, res, err = newRequest("PUT", uri, body, opt); err != nil {
 		return
 		return
 	}
 	}
-	defer res.Body.Close()
-	msg, err = readBody(res)
+
+	msg, err = readBody(req, res)
 	return
 	return
 }
 }
 
 
 // Delete HTTP request DELETE
 // Delete HTTP request DELETE
 func Delete(uri string, body io.Reader, opt RequestOption) (msg Message, err error) {
 func Delete(uri string, body io.Reader, opt RequestOption) (msg Message, err error) {
-	var res *http.Response
-	if res, err = newRequest("DELETE", uri, body, opt); err != nil {
+	var (
+		res *http.Response
+		req *http.Request
+	)
+	if req, res, err = newRequest("DELETE", uri, body, opt); err != nil {
 		return
 		return
 	}
 	}
-	defer res.Body.Close()
-	msg, err = readBody(res)
+
+	msg, err = readBody(req, res)
 	return
 	return
 }
 }
 
 
 // Connect HTTP request CONNECT
 // Connect HTTP request CONNECT
 func Connect(uri string, body io.Reader, opt RequestOption) (msg Message, err error) {
 func Connect(uri string, body io.Reader, opt RequestOption) (msg Message, err error) {
-	var res *http.Response
-	if res, err = newRequest("CONNECT", uri, body, opt); err != nil {
+	var (
+		res *http.Response
+		req *http.Request
+	)
+	if req, res, err = newRequest("CONNECT", uri, body, opt); err != nil {
 		return
 		return
 	}
 	}
-	defer res.Body.Close()
-	msg, err = readBody(res)
+
+	msg, err = readBody(req, res)
 	return
 	return
 }
 }
 
 
 // Options HTTP request OPTIONS
 // Options HTTP request OPTIONS
 func Options(uri string, body io.Reader, opt RequestOption) (msg Message, err error) {
 func Options(uri string, body io.Reader, opt RequestOption) (msg Message, err error) {
-	var res *http.Response
-	if res, err = newRequest("OPTIONS", uri, body, opt); err != nil {
+	var (
+		res *http.Response
+		req *http.Request
+	)
+	if req, res, err = newRequest("OPTIONS", uri, body, opt); err != nil {
 		return
 		return
 	}
 	}
-	defer res.Body.Close()
-	msg, err = readBody(res)
+
+	msg, err = readBody(req, res)
 	return
 	return
 }
 }
 
 
 // Trace HTTP request TRACE
 // Trace HTTP request TRACE
 func Trace(uri string, body io.Reader, opt RequestOption) (msg Message, err error) {
 func Trace(uri string, body io.Reader, opt RequestOption) (msg Message, err error) {
-	var res *http.Response
-	if res, err = newRequest("TRACE", uri, body, opt); err != nil {
+	var (
+		res *http.Response
+		req *http.Request
+	)
+	if req, res, err = newRequest("TRACE", uri, body, opt); err != nil {
 		return
 		return
 	}
 	}
-	defer res.Body.Close()
-	msg, err = readBody(res)
+
+	msg, err = readBody(req, res)
 	return
 	return
 }
 }
 
 
 // Patch HTTP request PATCH
 // Patch HTTP request PATCH
 func Patch(uri string, body io.Reader, opt RequestOption) (msg Message, err error) {
 func Patch(uri string, body io.Reader, opt RequestOption) (msg Message, err error) {
-	var res *http.Response
-	if res, err = newRequest("PATCH", uri, body, opt); err != nil {
+	var (
+		res *http.Response
+		req *http.Request
+	)
+	if req, res, err = newRequest("PATCH", uri, body, opt); err != nil {
 		return
 		return
 	}
 	}
-	defer res.Body.Close()
-	msg, err = readBody(res)
+
+	msg, err = readBody(req, res)
 	return
 	return
 }
 }

+ 17 - 0
xhttp/option.go

@@ -1,6 +1,7 @@
 package xhttp
 package xhttp
 
 
 import (
 import (
+	"net/http"
 	"time"
 	"time"
 )
 )
 
 
@@ -10,6 +11,9 @@ type RequestOption struct {
 	ResponseHeaderTimeout time.Duration
 	ResponseHeaderTimeout time.Duration
 
 
 	Headers map[string]string
 	Headers map[string]string
+	Cookies []*http.Cookie
+
+	username, password string
 
 
 	CertFile string
 	CertFile string
 	KeyFile  string
 	KeyFile  string
@@ -76,6 +80,19 @@ func (r *RequestOption) SetHeader(k, v string) *RequestOption {
 	return r
 	return r
 }
 }
 
 
+// AddCookie adds a cookie to the request.
+func (r *RequestOption) AddCookie(c *http.Cookie) *RequestOption {
+	r.Cookies = append(r.Cookies, c)
+	return r
+}
+
+// SetBasicAuth sets the request's Authorization header to use HTTP Basic Authentication with the provided username and password.
+func (r *RequestOption) SetBasicAuth(username, password string) *RequestOption {
+	r.username = username
+	r.password = password
+	return r
+}
+
 // AcceptEncodingGZIP Accept-Encoding gzip
 // AcceptEncodingGZIP Accept-Encoding gzip
 func (r *RequestOption) AcceptEncodingGZIP() *RequestOption {
 func (r *RequestOption) AcceptEncodingGZIP() *RequestOption {
 	r.Headers[AcceptEncoding] = "gzip, deflate, br"
 	r.Headers[AcceptEncoding] = "gzip, deflate, br"

+ 24 - 0
xhttp/response.go

@@ -4,10 +4,14 @@ import (
 	"encoding/json"
 	"encoding/json"
 	"encoding/xml"
 	"encoding/xml"
 	"net/http"
 	"net/http"
+	"net/url"
 )
 )
 
 
 // Message HTTP response
 // Message HTTP response
 type Message struct {
 type Message struct {
+	req *http.Request
+	res *http.Response
+
 	StatusCode int
 	StatusCode int
 	Body       []byte
 	Body       []byte
 	Header     http.Header
 	Header     http.Header
@@ -22,3 +26,23 @@ func (m Message) JSON(v interface{}) error {
 func (m Message) XML(v interface{}) error {
 func (m Message) XML(v interface{}) error {
 	return xml.Unmarshal(m.Body, v)
 	return xml.Unmarshal(m.Body, v)
 }
 }
+
+// Cookies parses and returns the cookies set in the Set-Cookie headers.
+func (m Message) Cookies() []*http.Cookie {
+	return m.res.Cookies()
+}
+
+// Location returns the URL of the response's "Location" header, if present. Relative redirects are resolved relative to the Response's Request. ErrNoLocation is returned if no Location header is present.
+func (m Message) Location() (*url.URL, error) {
+	return m.res.Location()
+}
+
+// UserAgent returns the client's User-Agent, if sent in the request.
+func (m Message) UserAgent() string {
+	return m.req.UserAgent()
+}
+
+// Referer returns the referring URL, if sent in the request.
+func (m Message) Referer() string {
+	return m.req.Referer()
+}

+ 1 - 1
xhttp/serve.go

@@ -24,7 +24,7 @@ func run(addr string, router http.Handler, option ServerOption) (errs, err error
 	errc := make(chan error)
 	errc := make(chan error)
 	go func() {
 	go func() {
 		c := make(chan os.Signal, 1)
 		c := make(chan os.Signal, 1)
-		signal.Notify(c, syscall.SIGHUP, syscall.SIGINT, syscall.SIGQUIT, syscall.SIGTERM)
+		signal.Notify(c, syscall.SIGHUP, syscall.SIGINT, syscall.SIGQUIT, syscall.SIGTERM, os.Interrupt)
 		errc <- fmt.Errorf("%s", <-c)
 		errc <- fmt.Errorf("%s", <-c)
 	}()
 	}()