spark-server/internal/transport/tcp/dispatch.go

package tcp

import (
	"context"
	"strconv"

	"go.uber.org/zap"

	"spark-server/internal/infrastructure/logger"
	"spark-server/internal/pkg/errs"
	"spark-server/internal/pkg/i18n"
	"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.CmdAuth {
		if !s.allowAuthIP(st.sess.RealIP()) {
			s.metrics.Inc("tcp_auth_rate_limited_total")
			s.replyErr(st, frame, errs.ErrRateLimited)
			return
		}
	} else 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, protocol.EncodeOK(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) {
	rc := st.sess.RequestContext()
	if frame.RequestID != 0 {
		rc.RequestID = strconv.FormatUint(frame.RequestID, 10)
	} else {
		rc.RequestID = reqctx.NewID()
	}
	ctx := logger.Attach(reqctx.With(context.Background(), rc), s.log)
	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)
	var respBody []byte
	if err != nil {
		s.logAppErr(ctx, err)
		respBody = protocol.EncodeErrLang(i18n.Parse(conn.Language()), err)
		_ = conn.Send(&session.Packet{
			Command: frame.Command, SequenceID: frame.SequenceID, RequestID: frame.RequestID, Body: respBody,
		})
	} else if resp != nil {
		respBody = resp.Body
		_ = conn.Send(&session.Packet{
			Command: resp.Command, SequenceID: resp.SequenceID, RequestID: resp.RequestID, Flags: resp.Flags, Body: resp.Body,
		})
	}
	logger.From(ctx).Info("tcp_access",
		zap.Uint16("cmd", frame.Command),
		zap.Uint64("frame_request_id", frame.RequestID),
		zap.String("request_body", logger.ClipBody(frame.Body)),
		zap.String("response_body", logger.ClipBody(respBody)),
	)
}

func (s *Server) replyErr(st *connState, frame protocol.Frame, err error) {
	rc := st.sess.RequestContext()
	if frame.RequestID != 0 {
		rc.RequestID = strconv.FormatUint(frame.RequestID, 10)
	}
	s.logAppErr(logger.Attach(reqctx.With(context.Background(), rc), s.log), err)
	s.writeEncoded(st, replyFrame(frame.Command, frame.SequenceID, frame.RequestID, protocol.EncodeErrLang(i18n.Parse(st.sess.Language()), err), 0))
}

func (s *Server) logAppErr(ctx context.Context, err error) {
	app, ok := errs.AsAppError(err)
	if !ok || app == nil || !app.LogPrint {
		return
	}
	logger.From(ctx).Error("app_error",
		zap.String("code", app.Code),
		zap.String("msg", app.Msg),
		zap.String("inner_msg", app.InnerMsg),
	)
}

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)
}