Reject invalid connection

This commit is contained in:
世界
2023-09-07 09:09:03 +08:00
parent 4ea2d460f4
commit 6b943caf37
9 changed files with 93 additions and 60 deletions

View File

@@ -150,7 +150,10 @@ func (c *Conn) Write(b []byte) (n int, err error) {
func (c *Conn) WriteBuffer(buffer *buf.Buffer) error {
if !c.requestWritten {
EncodeRequest(c.request, buf.With(buffer.ExtendHeader(RequestLen(c.request))))
err := EncodeRequest(c.request, buf.With(buffer.ExtendHeader(RequestLen(c.request))))
if err != nil {
return err
}
c.requestWritten = true
}
return c.ExtendedConn.WriteBuffer(buffer)
@@ -159,7 +162,11 @@ func (c *Conn) WriteBuffer(buffer *buf.Buffer) error {
func (c *Conn) WriteVectorised(buffers []*buf.Buffer) error {
if !c.requestWritten {
buffer := buf.NewSize(RequestLen(c.request))
EncodeRequest(c.request, buffer)
err := EncodeRequest(c.request, buffer)
if err != nil {
buffer.Release()
return err
}
c.requestWritten = true
return c.writer.WriteVectorised(append([]*buf.Buffer{buffer}, buffers...))
}