98 lines
2.1 KiB
Go
98 lines
2.1 KiB
Go
package tcp
|
|
|
|
import (
|
|
"errors"
|
|
"io"
|
|
"net"
|
|
"strconv"
|
|
|
|
"github.com/panjf2000/gnet/v2"
|
|
|
|
"spark-server/internal/transport/tcp/protocol"
|
|
)
|
|
|
|
func (s *Server) consumeProxy(st *connState) gnet.Action {
|
|
if !s.proxyProtocol || st.sess.ProxyDone() {
|
|
return gnet.None
|
|
}
|
|
avail := st.gc.InboundBuffered()
|
|
if avail == 0 {
|
|
return gnet.None
|
|
}
|
|
buf, err := st.gc.Peek(avail)
|
|
if err != nil {
|
|
return gnet.Close
|
|
}
|
|
h, n, err := protocol.ParsePrefix(buf)
|
|
if protocol.NeedMore(err) {
|
|
return gnet.None
|
|
}
|
|
if err != nil {
|
|
return proxyFallback(st, buf)
|
|
}
|
|
_, _ = st.gc.Discard(n)
|
|
st.sess.SetRealIP(headerRealIP(h, st.sess.RemoteAddr()))
|
|
return gnet.None
|
|
}
|
|
|
|
func proxyFallback(st *connState, buf []byte) gnet.Action {
|
|
if protocol.LooksLikeFrame(buf) {
|
|
st.sess.MarkProxySkipped()
|
|
return gnet.None
|
|
}
|
|
if protocol.MaybeFramePrefix(buf) && len(buf) < 2 {
|
|
return gnet.None
|
|
}
|
|
return gnet.Close
|
|
}
|
|
|
|
func headerRealIP(h protocol.Header, fallback string) string {
|
|
if h.SrcIP == nil {
|
|
return fallback
|
|
}
|
|
return net.JoinHostPort(h.SrcIP.String(), strconv.Itoa(h.SrcPort))
|
|
}
|
|
|
|
func (s *Server) consumeFrame(st *connState) (gnet.Action, bool) {
|
|
if st.gc.InboundBuffered() < protocol.HeaderSize {
|
|
return gnet.None, false
|
|
}
|
|
header, err := st.gc.Peek(protocol.HeaderSize)
|
|
if err != nil {
|
|
if errors.Is(err, io.ErrShortBuffer) {
|
|
return gnet.None, false
|
|
}
|
|
return gnet.Close, false
|
|
}
|
|
bodyLen, err := protocol.PeekBodyLength(header)
|
|
if err != nil || bodyLen > s.maxFrame {
|
|
return gnet.Close, false
|
|
}
|
|
total := protocol.HeaderSize + bodyLen
|
|
if st.gc.InboundBuffered() < total {
|
|
return gnet.None, false
|
|
}
|
|
raw, err := st.gc.Peek(total)
|
|
if err != nil {
|
|
if errors.Is(err, io.ErrShortBuffer) {
|
|
return gnet.None, false
|
|
}
|
|
return gnet.Close, false
|
|
}
|
|
frame, err := protocol.Decode(raw, s.maxFrame)
|
|
if err != nil {
|
|
return gnet.Close, false
|
|
}
|
|
_, _ = st.gc.Discard(total)
|
|
body, err := protocol.DecodeBody(frame.Body, frame.Flags, s.maxFrame)
|
|
if err != nil {
|
|
return gnet.Close, false
|
|
}
|
|
frame.Body = body
|
|
if frame.Version != 0 && frame.Version != protocol.CurrentVersion {
|
|
return gnet.Close, false
|
|
}
|
|
s.handleInbound(st, frame)
|
|
return gnet.None, true
|
|
}
|