Clash.Meta/common/callback/callback.go

78 lines
1.4 KiB
Go
Raw Normal View History

package callback
import (
"github.com/Dreamacro/clash/common/buf"
N "github.com/Dreamacro/clash/common/net"
C "github.com/Dreamacro/clash/constant"
)
type firstWriteCallBackConn struct {
C.Conn
callback func(error)
written bool
}
func (c *firstWriteCallBackConn) Write(b []byte) (n int, err error) {
defer func() {
if !c.written {
c.written = true
c.callback(err)
}
}()
return c.Conn.Write(b)
}
func (c *firstWriteCallBackConn) Upstream() any {
return c.Conn
}
type extendedConn interface {
C.Conn
N.ExtendedConn
}
type firstWriteCallBackExtendedConn struct {
extendedConn
callback func(error)
written bool
}
func (c *firstWriteCallBackExtendedConn) Write(b []byte) (n int, err error) {
defer func() {
if !c.written {
c.written = true
c.callback(err)
}
}()
return c.extendedConn.Write(b)
}
func (c *firstWriteCallBackExtendedConn) WriteBuffer(buffer *buf.Buffer) (err error) {
defer func() {
if !c.written {
c.written = true
c.callback(err)
}
}()
return c.extendedConn.WriteBuffer(buffer)
}
func (c *firstWriteCallBackExtendedConn) Upstream() any {
return c.extendedConn
}
func NewFirstWriteCallBackConn(c C.Conn, callback func(error)) C.Conn {
if c, ok := c.(extendedConn); ok {
return &firstWriteCallBackExtendedConn{
extendedConn: c,
callback: callback,
written: false,
}
}
return &firstWriteCallBackConn{
Conn: c,
callback: callback,
written: false,
}
}