Skip to content

Commit

Permalink
Update.
Browse files Browse the repository at this point in the history
  • Loading branch information
xOS committed Aug 25, 2024
1 parent 84c7d7c commit e35a11c
Showing 1 changed file with 56 additions and 0 deletions.
56 changes: 56 additions & 0 deletions pkg/utils/request_wrapper.go
Original file line number Diff line number Diff line change
@@ -0,0 +1,56 @@
package utils

import (
"bytes"
"io"
"net"
"net/http"

"github.com/gin-gonic/gin"
)

var _ io.ReadWriteCloser = &RequestWrapper{}

type RequestWrapper struct {
req *http.Request
reader *bytes.Buffer
writer net.Conn
}

func NewRequestWrapper(req *http.Request, writer gin.ResponseWriter) (*RequestWrapper, error) {
conn, _, err := writer.Hijack()
if err != nil {
return nil, err
}
buf := bytes.NewBuffer(nil)
if err = req.Write(buf); err != nil {
return nil, err
}
return &RequestWrapper{
req: req,
reader: buf,
writer: conn,
}, nil
}

func (rw *RequestWrapper) Read(p []byte) (int, error) {
count, err := rw.reader.Read(p)
if err == nil {
return count, nil
}
if err != io.EOF {
return count, err
}
// request 数据读完之后等待客户端断开连接或 grpc 超时
return rw.writer.Read(p)
}

func (rw *RequestWrapper) Write(p []byte) (int, error) {
return rw.writer.Write(p)
}

func (rw *RequestWrapper) Close() error {
rw.req.Body.Close()
rw.writer.Close()
return nil
}

0 comments on commit e35a11c

Please sign in to comment.