package rcon

import "time"

// Settings contains option to Conn.
type Settings struct {
        dialTimeout   time.Duration
        deadline      time.Duration
        maxCommandLen int
}

// DefaultSettings provides default deadline settings to Conn.
var DefaultSettings = Settings{
        dialTimeout:   DefaultDialTimeout,
        deadline:      DefaultDeadline,
        maxCommandLen: DefaultMaxCommandLen,
}

// Option allows to inject settings to Settings.
type Option func(s *Settings)

// SetDialTimeout injects dial Timeout to Settings.
func SetDialTimeout(timeout time.Duration) Option {
        return func(s *Settings) {
                s.dialTimeout = timeout
        }
}

// SetDeadline injects read/write Timeout to Settings.
func SetDeadline(timeout time.Duration) Option {
        return func(s *Settings) {
                s.deadline = timeout
        }
}

// SetMaxCommandLen injects maxCommandLen to Settings.
func SetMaxCommandLen(maxCommandLen int) Option {
        return func(s *Settings) {
                s.maxCommandLen = maxCommandLen
        }
}

package rcon

import (
        "bytes"
        "encoding/binary"
        "fmt"
        "io"
)

// Packet sizes definitions.
const (
        PacketPaddingSize int32 = 2 // Size of Packet's padding.
        PacketHeaderSize  int32 = 8 // Size of Packet's header.

        MinPacketSize = PacketPaddingSize + PacketHeaderSize
        MaxPacketSize = 4096 + MinPacketSize
)

// Packet is a rcon packet. Both requests and responses are sent as
// TCP packets. Their payload follows the following basic structure.
type Packet struct {
        // The packet size field is a 32-bit little endian integer, representing
        // the length of the request in bytes. Note that the packet size field
        // itself is not included when determining the size of the packet,
        // so the value of this field is always 4 less than the packet's actual
        // length. The minimum possible value for packet size is 10.
        // The maximum possible value of packet size is 4096.
        // If the response is too large to fit into one packet, it will be split
        // and sent as multiple packets.
        Size int32

        // The packet id field is a 32-bit little endian integer chosen by the
        // client for each request. It may be set to any positive integer.
        // When the RemoteServer responds to the request, the response packet
        // will have the same packet id as the original request (unless it is
        // a failed SERVERDATA_AUTH_RESPONSE packet).
        // It need not be unique, but if a unique packet id is assigned,
        // it can be used to match incoming responses to their corresponding requests.
        ID int32

        // The packet type field is a 32-bit little endian integer, which indicates
        // the purpose of the packet. Its value will always be either 0, 2, or 3,
        // depending on which of the following request/response types the packet
        // represents:
        // SERVERDATA_AUTH = 3,
        // SERVERDATA_AUTH_RESPONSE = 2,
        // SERVERDATA_EXECCOMMAND = 2,
        // SERVERDATA_RESPONSE_VALUE = 0.
        Type int32

        // The packet body field is a null-terminated string encoded in ASCII
        // (i.e. ASCIIZ). Depending on the packet type, it may contain either the
        // RCON MockPassword for the RemoteServer, the command to be executed,
        // or the RemoteServer's response to a request.
        body []byte
}

// NewPacket creates and initializes a new Packet using packetType,
// packetID and body as its initial contents. NewPacket is intended to
// calculate packet size from body length and 10 bytes for rcon headers
// and termination strings.
func NewPacket(packetType int32, packetID int32, body string) *Packet {
        size := len([]byte(body)) + int(PacketHeaderSize+PacketPaddingSize)

        return &Packet{
                Size: int32(size), //nolint:gosec // No matter
                Type: packetType,
                ID:   packetID,
                body: []byte(body),
        }
}

// Body returns packet bytes body as a string.
func (packet *Packet) Body() string {
        return string(packet.body)
}

// WriteTo implements io.WriterTo for write a packet to w.
func (packet *Packet) WriteTo(w io.Writer) (int64, error) {
        buffer := bytes.NewBuffer(make([]byte, 0, packet.Size+4))

        _ = binary.Write(buffer, binary.LittleEndian, packet.Size)
        _ = binary.Write(buffer, binary.LittleEndian, packet.ID)
        _ = binary.Write(buffer, binary.LittleEndian, packet.Type)

        // Write command body, null terminated ASCII string and an empty ASCIIZ string.
        buffer.Write(append(packet.body, 0x00, 0x00))

        return buffer.WriteTo(w)
}

// ReadFrom implements io.ReaderFrom for read a packet from r.
func (packet *Packet) ReadFrom(r io.Reader) (int64, error) {
        var n int64

        if err := binary.Read(r, binary.LittleEndian, &packet.Size); err != nil {
                return n, fmt.Errorf("rcon: read packet size: %w", err)
        }

        n += 4

        if packet.Size < MinPacketSize {
                return n, ErrResponseTooSmall
        }

        if err := binary.Read(r, binary.LittleEndian, &packet.ID); err != nil {
                return n, fmt.Errorf("rcon: read packet id: %w", err)
        }

        n += 4

        if err := binary.Read(r, binary.LittleEndian, &packet.Type); err != nil {
                return n, fmt.Errorf("rcon: read packet type: %w", err)
        }

        n += 4

        // String can actually include null characters which is the case in
        // response to a SERVERDATA_RESPONSE_VALUE packet.
        packet.body = make([]byte, packet.Size-PacketHeaderSize)

        var i int64
        for i < int64(packet.Size-PacketHeaderSize) {
                var m int
                var err error

                if m, err = r.Read(packet.body[i:]); err != nil {
                        return n + int64(m) + i, fmt.Errorf("rcon: %w", err)
                }

                i += int64(m)
        }

        n += i

        // Remove null terminated strings from response body.
        if !bytes.Equal(packet.body[len(packet.body)-int(PacketPaddingSize):], []byte{0x00, 0x00}) {
                return n, ErrInvalidPacketPadding
        }

        packet.body = packet.body[0 : len(packet.body)-int(PacketPaddingSize)]

        return n, nil
}

// Package rcon implements Source RCON Protocol which is described in the
// documentation: https://developer.valvesoftware.com/wiki/Source_RCON_Protocol.
package rcon

import (
        "encoding/binary"
        "errors"
        "fmt"
        "net"
        "time"
)

const (
        // DefaultDialTimeout provides default auth timeout to remote server.
        DefaultDialTimeout = 5 * time.Second

        // DefaultDeadline provides default deadline to tcp read/write operations.
        DefaultDeadline = 5 * time.Second

        // DefaultMaxCommandLen is an artificial restriction, but it will help in case of random
        // large queries.
        DefaultMaxCommandLen = 1000

        // SERVERDATA_AUTH is the first packet sent by the client,
        // which is used to authenticate the conn with the server.
        SERVERDATA_AUTH int32 = 3

        // SERVERDATA_AUTH_ID is any positive integer, chosen by the client
        // (will be mirrored back in the server's response).
        SERVERDATA_AUTH_ID int32 = 0

        // SERVERDATA_AUTH_RESPONSE packet is a notification of the conn's current auth
        // status. When the server receives an auth request, it will respond with an empty
        // SERVERDATA_RESPONSE_VALUE, followed immediately by a SERVERDATA_AUTH_RESPONSE
        // indicating whether authentication succeeded or failed. Note that the status
        // code is returned in the packet id field, so when pairing the response with
        // the original auth request, you may need to look at the packet id of the
        // preceding SERVERDATA_RESPONSE_VALUE.
        // If authentication was successful, the ID assigned by the request.
        // If auth failed, -1 (0xFF FF FF FF).
        SERVERDATA_AUTH_RESPONSE int32 = 2

        // SERVERDATA_RESPONSE_VALUE packet is the response to a SERVERDATA_EXECCOMMAND
        // request. The ID assigned by the original request.
        SERVERDATA_RESPONSE_VALUE int32 = 0

        // SERVERDATA_EXECCOMMAND packet type represents a command issued to the server
        // by a client. The response will vary depending on the command issued.
        SERVERDATA_EXECCOMMAND int32 = 2

        // SERVERDATA_EXECCOMMAND_ID is any positive integer, chosen by the client
        // (will be mirrored back in the server's response).
        SERVERDATA_EXECCOMMAND_ID int32 = 0
)

var (
        // ErrAuthNotRCON is returned when got auth response with negative size.
        ErrAuthNotRCON = errors.New("response from not rcon server")

        // ErrInvalidAuthResponse is returned when we didn't get an auth packet
        // back for second read try after discard empty SERVERDATA_RESPONSE_VALUE
        // from authentication response.
        ErrInvalidAuthResponse = errors.New("invalid authentication packet type response")

        // ErrAuthFailed is returned when the package id from authentication
        // response is -1.
        ErrAuthFailed = errors.New("authentication failed")

        // ErrInvalidPacketID is returned when the package id from server response
        // was not mirrored back from request.
        ErrInvalidPacketID = errors.New("response for another request")

        // ErrInvalidPacketPadding is returned when the bytes after type field from
        // response is not equal to null-terminated ASCII strings.
        ErrInvalidPacketPadding = errors.New("invalid response padding")

        // ErrResponseTooSmall is returned when the server response is smaller
        // than 10 bytes.
        ErrResponseTooSmall = errors.New("response too small")

        // ErrCommandTooLong is returned when executed command length is bigger
        // than MaxCommandLen characters.
        ErrCommandTooLong = errors.New("command too long")

        // ErrCommandEmpty is returned when executed command length equal 0.
        ErrCommandEmpty = errors.New("command too small")

        // ErrMultiErrorOccurred is returned when close connection failed with
        // error after auth failed.
        ErrMultiErrorOccurred = errors.New("an error occurred while handling another error")
)

// Conn is source RCON generic stream-oriented network connection.
type Conn struct {
        conn     net.Conn
        settings Settings
}

// open creates a new Conn from an existing net.Conn and authenticates it.
func open(conn net.Conn, password string, settings Settings) (*Conn, error) {
        client := Conn{conn: conn, settings: settings}

        if err := client.auth(password); err != nil {
                // Failed to auth conn with the server.
                if err2 := client.Close(); err2 != nil {
                        return &client, fmt.Errorf("%w: %s. Previous error: %s", ErrMultiErrorOccurred, err2.Error(), err.Error())
                }

                return &client, fmt.Errorf("rcon: %w", err)
        }

        return &client, nil
}

// Open creates a new authorized Conn from an existing net.Conn.
func Open(conn net.Conn, password string, options ...Option) (*Conn, error) {
        settings := DefaultSettings
        for _, option := range options {
                option(&settings)
        }

        return open(conn, password, settings)
}

// Dial creates a new authorized Conn tcp dialer connection.
func Dial(address string, password string, options ...Option) (*Conn, error) {
        settings := DefaultSettings

        for _, option := range options {
                option(&settings)
        }

        conn, err := net.DialTimeout("tcp", address, settings.dialTimeout)
        if err != nil {
                // Failed to open TCP connection to the server.
                return nil, fmt.Errorf("rcon: %w", err)
        }

        return open(conn, password, settings)
}

// Execute sends command type and it string to execute to the remote server,
// creating a packet with a SERVERDATA_EXECCOMMAND_ID for the server to mirror,
// and compiling its payload bytes in the appropriate order. The response body
// is decompiled from bytes into a string for return.
func (c *Conn) Execute(command string) (string, error) {
        if command == "" {
                return "", ErrCommandEmpty
        }

        if c.settings.maxCommandLen > 0 && len(command) > c.settings.maxCommandLen {
                return "", ErrCommandTooLong
        }

        if err := c.write(SERVERDATA_EXECCOMMAND, SERVERDATA_EXECCOMMAND_ID, command); err != nil {
                return "", err
        }

        response, err := c.read()
        if err != nil {
                return response.Body(), err
        }

        if response.ID != SERVERDATA_EXECCOMMAND_ID {
                return response.Body(), ErrInvalidPacketID
        }

        return response.Body(), nil
}

// LocalAddr returns the local network address.
func (c *Conn) LocalAddr() net.Addr {
        return c.conn.LocalAddr()
}

// RemoteAddr returns the remote network address.
func (c *Conn) RemoteAddr() net.Addr {
        return c.conn.RemoteAddr()
}

// Close closes the connection.
func (c *Conn) Close() error {
        return c.conn.Close()
}

// auth sends SERVERDATA_AUTH request to the remote server and
// authenticates client for the next requests.
func (c *Conn) auth(password string) error {
        if err := c.write(SERVERDATA_AUTH, SERVERDATA_AUTH_ID, password); err != nil {
                return err
        }

        if c.settings.deadline != 0 {
                if err := c.conn.SetReadDeadline(time.Now().Add(c.settings.deadline)); err != nil {
                        return fmt.Errorf("rcon: %w", err)
                }
        }

        response, err := c.readHeader()
        if err != nil {
                return err
        }

        size := response.Size - PacketHeaderSize
        if size < 0 {
                return ErrAuthNotRCON
        }

        // When the server receives an auth request, it will respond with an empty
        // SERVERDATA_RESPONSE_VALUE, followed immediately by a SERVERDATA_AUTH_RESPONSE
        // indicating whether authentication succeeded or failed.
        // Some servers doesn't send an empty SERVERDATA_RESPONSE_VALUE packet, so we
        // do this case optional.
        if response.Type == SERVERDATA_RESPONSE_VALUE {
                // Discard empty SERVERDATA_RESPONSE_VALUE from authentication response.
                _, _ = c.conn.Read(make([]byte, size))

                if response, err = c.readHeader(); err != nil {
                        return err
                }
        }

        // We must to read response body.
        buffer := make([]byte, size)
        if _, err := c.conn.Read(buffer); err != nil {
                return fmt.Errorf("rcon: %w", err)
        }

        if response.Type != SERVERDATA_AUTH_RESPONSE {
                return ErrInvalidAuthResponse
        }

        if response.ID == -1 {
                return ErrAuthFailed
        }

        if response.ID != SERVERDATA_AUTH_ID {
                return ErrInvalidPacketID
        }

        return nil
}

// write creates packet and writes it to established tcp conn.
func (c *Conn) write(packetType int32, packetID int32, command string) error {
        if c.settings.deadline != 0 {
                if err := c.conn.SetWriteDeadline(time.Now().Add(c.settings.deadline)); err != nil {
                        return fmt.Errorf("rcon: %w", err)
                }
        }

        packet := NewPacket(packetType, packetID, command)
        _, err := packet.WriteTo(c.conn)

        return err
}

// read reads structured binary data from c.conn into packet.
func (c *Conn) read() (*Packet, error) {
        if c.settings.deadline != 0 {
                if err := c.conn.SetReadDeadline(time.Now().Add(c.settings.deadline)); err != nil {
                        return nil, fmt.Errorf("rcon: %w", err)
                }
        }

        packet := &Packet{}
        if _, err := packet.ReadFrom(c.conn); err != nil {
                return packet, err
        }

        // Workaround for Rust server.
        // Rust rcon server responses packet with a type of 4 and the next packet
        // is valid. It is undocumented, so skip packet and read next.
        if packet.Type == 4 {
                if _, err := packet.ReadFrom(c.conn); err != nil {
                        return packet, err
                }

                // One more workaround for Rust server.
                // When sent command "Say" there is no response data from server with
                // packet.ID = SERVERDATA_EXECCOMMAND_ID, only previous console message
                // that command was received with packet.ID = -1, therefore, forcibly
                // set packet.ID to SERVERDATA_EXECCOMMAND_ID.
                if packet.ID == -1 {
                        packet.ID = SERVERDATA_EXECCOMMAND_ID
                }
        }

        return packet, nil
}

// readHeader reads structured binary data without body from c.conn into packet.
func (c *Conn) readHeader() (Packet, error) {
        var packet Packet
        if err := binary.Read(c.conn, binary.LittleEndian, &packet.Size); err != nil {
                return packet, fmt.Errorf("rcon: read packet size: %w", err)
        }

        if err := binary.Read(c.conn, binary.LittleEndian, &packet.ID); err != nil {
                return packet, fmt.Errorf("rcon: read packet id: %w", err)
        }

        if err := binary.Read(c.conn, binary.LittleEndian, &packet.Type); err != nil {
                return packet, fmt.Errorf("rcon: read packet type: %w", err)
        }

        return packet, nil
}

package rcontest

import (
        "net"

        "github.com/gorcon/rcon"
)

// Context represents the context of the current RCON request. It holds request
// and conn objects.
type Context struct {
        server  *Server
        conn    net.Conn
        request *rcon.Packet
}

// Server returns the Server instance.
func (c *Context) Server() *Server {
        return c.server
}

// Conn returns current RCON connection.
func (c *Context) Conn() net.Conn {
        return c.conn
}

// Request returns received *rcon.Packet.
func (c *Context) Request() *rcon.Packet {
        return c.request
}

package rcontest

// Option allows to inject Settings to Server.
type Option func(s *Server)

// SetSettings injects configuration for RCON Server.
func SetSettings(settings Settings) Option {
        return func(s *Server) {
                s.Settings = settings
        }
}

// SetAuthHandler injects HandlerFunc with authorisation data checking.
func SetAuthHandler(handler HandlerFunc) Option {
        return func(s *Server) {
                s.SetAuthHandler(handler)
        }
}

// SetCommandHandler injects HandlerFunc with commands processing.
func SetCommandHandler(handler HandlerFunc) Option {
        return func(s *Server) {
                s.SetCommandHandler(handler)
        }
}

// Package rcontest contains RCON server for RCON client testing.
package rcontest

import (
        "errors"
        "fmt"
        "io"
        "net"
        "sync"
        "time"

        "github.com/gorcon/rcon"
)

// Server is an RCON server listening on a system-chosen port on the
// local loopback interface, for use in end-to-end RCON tests.
type Server struct {
        Settings       Settings
        Listener       net.Listener
        addr           string
        authHandler    HandlerFunc
        commandHandler HandlerFunc
        connections    map[net.Conn]struct{}
        quit           chan bool
        wg             sync.WaitGroup
        mu             sync.Mutex
        closed         bool
}

// Settings contains configuration for RCON Server.
type Settings struct {
        Password             string
        AuthResponseDelay    time.Duration
        CommandResponseDelay time.Duration
}

// HandlerFunc defines a function to serve RCON requests.
type HandlerFunc func(c *Context)

// AuthHandler checks authorisation data and responses with
// SERVERDATA_AUTH_RESPONSE packet.
func AuthHandler(c *Context) {
        if c.Request().Body() == c.Server().Settings.Password {
                // First write SERVERDATA_RESPONSE_VALUE packet with empty body.
                _, _ = rcon.NewPacket(rcon.SERVERDATA_RESPONSE_VALUE, c.Request().ID, "").WriteTo(c.Conn())

                // Than write SERVERDATA_AUTH_RESPONSE packet to allow authHandler success.
                _, _ = rcon.NewPacket(rcon.SERVERDATA_AUTH_RESPONSE, rcon.SERVERDATA_AUTH_ID, "").WriteTo(c.Conn())
        } else {
                // If authentication was failed, the ID must be assigned to -1.
                _, _ = rcon.NewPacket(rcon.SERVERDATA_AUTH_RESPONSE, -1, string([]byte{0x00})).WriteTo(c.Conn())
        }
}

// EmptyHandler responses with empty body. Is used when start RCON Server with nil
// commandHandler.
func EmptyHandler(c *Context) {
        _, _ = rcon.NewPacket(rcon.SERVERDATA_RESPONSE_VALUE, c.Request().ID, "").WriteTo(c.Conn())
}

func newLocalListener() net.Listener {
        l, err := net.Listen("tcp", "127.0.0.1:0")
        if err != nil {
                panic(fmt.Sprintf("rcontest: failed to listen on a port: %v", err))
        }

        return l
}

// NewServer returns a running RCON Server or nil if an error occurred.
// The caller should call Close when finished, to shut it down.
func NewServer(options ...Option) *Server {
        server := NewUnstartedServer(options...)
        server.Start()

        return server
}

// NewUnstartedServer returns a new Server but doesn't start it.
// After changing its configuration, the caller should call Start.
// The caller should call Close when finished, to shut it down.
func NewUnstartedServer(options ...Option) *Server {
        server := Server{
                Listener:       newLocalListener(),
                authHandler:    AuthHandler,
                commandHandler: EmptyHandler,
                connections:    make(map[net.Conn]struct{}),
                quit:           make(chan bool),
        }

        for _, option := range options {
                option(&server)
        }

        return &server
}

// SetAuthHandler injects HandlerFunc with authorisation data checking.
func (s *Server) SetAuthHandler(handler HandlerFunc) {
        s.authHandler = handler
}

// SetCommandHandler injects HandlerFunc with commands processing.
func (s *Server) SetCommandHandler(handler HandlerFunc) {
        s.commandHandler = handler
}

// Start starts a server from NewUnstartedServer.
func (s *Server) Start() {
        if s.addr != "" {
                panic("server already started")
        }

        s.addr = s.Listener.Addr().String()
        s.goServe()
}

// Close shuts down the Server.
func (s *Server) Close() {
        if s.closed {
                return
        }

        s.closed = true
        close(s.quit)
        s.Listener.Close()

        // Waiting for server connections.
        s.wg.Wait()

        s.mu.Lock()
        for c := range s.connections {
                // Force-close any connections.
                s.closeConn(c)
        }
        s.mu.Unlock()
}

// Addr returns IPv4 string Server address.
func (s *Server) Addr() string {
        return s.addr
}

// NewContext returns a Context instance.
func (s *Server) NewContext(conn net.Conn) (*Context, error) {
        ctx := Context{server: s, conn: conn, request: &rcon.Packet{}}

        if _, err := ctx.request.ReadFrom(conn); err != nil {
                return &ctx, fmt.Errorf("rcontest: %w", err)
        }

        return &ctx, nil
}

// serve handles incoming requests until a stop signal is given with Close.
func (s *Server) serve() {
        for {
                conn, err := s.Listener.Accept()
                if err != nil {
                        if s.isRunning() {
                                panic(fmt.Errorf("rcontest: %w", err))
                        }

                        return
                }

                s.wg.Add(1)

                go s.handle(conn)
        }
}

// serve calls serve in goroutine.
func (s *Server) goServe() {
        s.wg.Add(1)

        go func() {
                defer s.wg.Done()

                s.serve()
        }()
}

// handle handles incoming client conn.
func (s *Server) handle(conn net.Conn) {
        s.mu.Lock()
        s.connections[conn] = struct{}{}
        s.mu.Unlock()

        defer func() {
                s.closeConn(conn)
                s.wg.Done()
        }()

        for {
                ctx, err := s.NewContext(conn)
                if err != nil {
                        if !errors.Is(err, io.EOF) {
                                panic(fmt.Errorf("failed read request: %w", err))
                        }

                        return
                }

                switch ctx.Request().Type {
                case rcon.SERVERDATA_AUTH:
                        if s.Settings.AuthResponseDelay != 0 {
                                time.Sleep(s.Settings.AuthResponseDelay)
                        }

                        s.authHandler(ctx)
                case rcon.SERVERDATA_EXECCOMMAND:
                        if s.Settings.CommandResponseDelay != 0 {
                                time.Sleep(s.Settings.CommandResponseDelay)
                        }

                        s.commandHandler(ctx)
                }
        }
}

// isRunning returns true if Server is running and false if is not.
func (s *Server) isRunning() bool {
        select {
        case <-s.quit:
                return false
        default:
                return true
        }
}

// closeConn closes a client conn and removes it from connections map.
func (s *Server) closeConn(conn net.Conn) {
        s.mu.Lock()
        defer s.mu.Unlock()

        if err := conn.Close(); err != nil {
                panic(fmt.Errorf("close conn error: %w", err))
        }

        delete(s.connections, conn)
}