Quickly close the peer connection use the tcp half-close method CloseWrite, instead of just relying on timeout

This commit is contained in:
ansoda 2021-02-10 17:08:07 +08:00
parent 42c15b9262
commit ab4b93f8b6
3 changed files with 54 additions and 0 deletions

View File

@ -55,10 +55,24 @@ func Relay(left, right net.Conn) error {
go func() { go func() {
defer wg.Done() defer wg.Done()
_, err1 = Copy(right, left) _, err1 = Copy(right, left)
if rc, ok := right.(interface {
CloseWrite() error
}); ok {
if e1 := rc.CloseWrite(); err1 == nil && e1 != io.EOF {
err1 = e1
}
}
right.SetReadDeadline(time.Now().Add(wait)) // unblock read on right right.SetReadDeadline(time.Now().Add(wait)) // unblock read on right
}() }()
_, err = Copy(left, right) _, err = Copy(left, right)
if lc, ok := left.(interface {
CloseWrite() error
}); ok {
if e := lc.CloseWrite(); err1 == nil && e != io.EOF {
err = e
}
}
left.SetReadDeadline(time.Now().Add(wait)) // unblock read on left left.SetReadDeadline(time.Now().Add(wait)) // unblock read on left
wg.Wait() wg.Wait()

View File

@ -65,3 +65,23 @@ func (c *streamConn) Write(b []byte) (int, error) {
} }
return c.w.Write(b) return c.w.Write(b)
} }
func (c *streamConn) CloseWrite() error {
if conn, ok := c.Conn.(interface {
CloseWrite() error
}); ok {
return conn.CloseWrite()
}
return nil
}
func (c *streamConn) CloseRead() error {
if conn, ok := c.Conn.(interface {
CloseRead() error
}); ok {
return conn.CloseRead()
}
return nil
}

View File

@ -60,3 +60,23 @@ func (c *conn) Write(b []byte) (int, error) {
} }
return c.w.Write(b) return c.w.Write(b)
} }
func (c *conn) CloseWrite() error {
if conn, ok := c.Conn.(interface {
CloseWrite() error
}); ok {
return conn.CloseWrite()
}
return nil
}
func (c *conn) CloseRead() error {
if conn, ok := c.Conn.(interface {
CloseRead() error
}); ok {
return conn.CloseRead()
}
return nil
}