94 lines
2.5 KiB
Go
94 lines
2.5 KiB
Go
package tcp
|
|
|
|
import (
|
|
"context"
|
|
|
|
"spark-server/internal/pkg/errs"
|
|
"spark-server/internal/pkg/reqctx"
|
|
"spark-server/internal/transport/tcp/protocol"
|
|
"spark-server/internal/transport/tcp/session"
|
|
)
|
|
|
|
func (s *Server) handleInbound(st *connState, frame protocol.Frame) {
|
|
st.sess.Touch()
|
|
if frame.Command != protocol.CmdHeartbeat && frame.Command != protocol.CmdPushAck && !st.limit.Allow() {
|
|
s.metrics.Inc("tcp_rate_limited_total")
|
|
s.replyErr(st, frame, errs.ErrRateLimited)
|
|
return
|
|
}
|
|
if err := ensureAuth(st.sess, frame.Command); err != nil {
|
|
s.replyErr(st, frame, err)
|
|
return
|
|
}
|
|
switch frame.Command {
|
|
case protocol.CmdHeartbeat:
|
|
s.writeEncoded(st, replyFrame(protocol.CmdHeartbeat, frame.SequenceID, frame.RequestID, nil, 0))
|
|
case protocol.CmdPushAck:
|
|
return
|
|
default:
|
|
s.submitBusiness(st, frame)
|
|
}
|
|
}
|
|
|
|
func ensureAuth(conn *session.Conn, cmd uint16) error {
|
|
switch cmd {
|
|
case protocol.CmdAuth, protocol.CmdHeartbeat, protocol.CmdPushAck:
|
|
return nil
|
|
default:
|
|
if !conn.Authenticated() {
|
|
return errs.ErrNotAuthenticated
|
|
}
|
|
return nil
|
|
}
|
|
}
|
|
|
|
func (s *Server) submitBusiness(st *connState, frame protocol.Frame) {
|
|
ctx := reqctx.With(context.Background(), st.sess.RequestContext())
|
|
copied := frame
|
|
err := s.pool.Submit(ctx, func() {
|
|
s.runHandler(ctx, st.sess, copied)
|
|
})
|
|
if err == errs.ErrBusy {
|
|
s.metrics.Inc("taskpool_rejected_total")
|
|
s.replyErr(st, frame, errs.ErrBusy)
|
|
}
|
|
}
|
|
|
|
func (s *Server) runHandler(ctx context.Context, conn *session.Conn, frame protocol.Frame) {
|
|
ctx, cancel := context.WithTimeout(ctx, s.handlerTimeout)
|
|
defer cancel()
|
|
resp, err := s.router.Dispatch(ctx, conn, frame)
|
|
if err != nil {
|
|
_ = conn.Send(&session.Packet{
|
|
Command: frame.Command, SequenceID: frame.SequenceID, RequestID: frame.RequestID,
|
|
Body: protocol.EncodeError(errorCode(err)),
|
|
})
|
|
return
|
|
}
|
|
if resp == nil {
|
|
return
|
|
}
|
|
_ = conn.Send(&session.Packet{
|
|
Command: resp.Command, SequenceID: resp.SequenceID, RequestID: resp.RequestID, Flags: resp.Flags, Body: resp.Body,
|
|
})
|
|
}
|
|
|
|
func (s *Server) replyErr(st *connState, frame protocol.Frame, err error) {
|
|
body := protocol.EncodeError(errorCode(err))
|
|
s.writeEncoded(st, replyFrame(frame.Command, frame.SequenceID, frame.RequestID, body, 0))
|
|
}
|
|
|
|
func (s *Server) writeEncoded(st *connState, frame protocol.Frame) {
|
|
body, flags, err := protocol.PrepareBody(frame.Body, s.compressMin)
|
|
if err != nil {
|
|
return
|
|
}
|
|
frame.Body = body
|
|
frame.Flags |= flags
|
|
raw, err := protocol.Encode(frame)
|
|
if err != nil {
|
|
return
|
|
}
|
|
_ = st.gc.AsyncWrite(raw, nil)
|
|
}
|