Skip to content

Commit e35a11c

Browse files
committed
Update.
1 parent 84c7d7c commit e35a11c

File tree

1 file changed

+56
-0
lines changed

1 file changed

+56
-0
lines changed

pkg/utils/request_wrapper.go

Lines changed: 56 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,56 @@
1+
package utils
2+
3+
import (
4+
"bytes"
5+
"io"
6+
"net"
7+
"net/http"
8+
9+
"github.com/gin-gonic/gin"
10+
)
11+
12+
var _ io.ReadWriteCloser = &RequestWrapper{}
13+
14+
type RequestWrapper struct {
15+
req *http.Request
16+
reader *bytes.Buffer
17+
writer net.Conn
18+
}
19+
20+
func NewRequestWrapper(req *http.Request, writer gin.ResponseWriter) (*RequestWrapper, error) {
21+
conn, _, err := writer.Hijack()
22+
if err != nil {
23+
return nil, err
24+
}
25+
buf := bytes.NewBuffer(nil)
26+
if err = req.Write(buf); err != nil {
27+
return nil, err
28+
}
29+
return &RequestWrapper{
30+
req: req,
31+
reader: buf,
32+
writer: conn,
33+
}, nil
34+
}
35+
36+
func (rw *RequestWrapper) Read(p []byte) (int, error) {
37+
count, err := rw.reader.Read(p)
38+
if err == nil {
39+
return count, nil
40+
}
41+
if err != io.EOF {
42+
return count, err
43+
}
44+
// request 数据读完之后等待客户端断开连接或 grpc 超时
45+
return rw.writer.Read(p)
46+
}
47+
48+
func (rw *RequestWrapper) Write(p []byte) (int, error) {
49+
return rw.writer.Write(p)
50+
}
51+
52+
func (rw *RequestWrapper) Close() error {
53+
rw.req.Body.Close()
54+
rw.writer.Close()
55+
return nil
56+
}

0 commit comments

Comments
 (0)