Remove MustNewServer function from rcontest, add NewUnstartedServer

gorcon · Nov 23, 2020 · f668773 · f668773
1 parent cdfa210
commit f668773
Show file tree

Hide file tree

Showing 4 changed files with 130 additions and 159 deletions.
diff --git a/rcon_test.go b/rcon_test.go
@@ -14,11 +14,11 @@ import (
 )
 
 func TestDial(t *testing.T) {
-	server := rcontest.MustNewServer()
-	defer server.MustClose()
+	server := rcontest.NewServer(nil)
+	defer server.Close()
 
 	t.Run("connection refused", func(t *testing.T) {
-		conn, err := rcon.Dial("127.0.0.2:12345", rcontest.MockPasswordRCON)
+		conn, err := rcon.Dial("127.0.0.2:12345", "password")
 		if !assert.Error(t, err) {
 			// Close connection if established.
 			assert.NoError(t, conn.Close())
@@ -43,19 +43,19 @@ func TestDial(t *testing.T) {
 	})
 
 	t.Run("auth success", func(t *testing.T) {
-		conn, err := rcon.Dial(server.Addr(), rcontest.MockPasswordRCON)
+		conn, err := rcon.Dial(server.Addr(), "password")
 		if assert.NoError(t, err) {
 			assert.NoError(t, conn.Close())
 		}
 	})
 }
 
 func TestConn_Execute(t *testing.T) {
-	server := rcontest.MustNewServer()
-	defer server.MustClose()
+	server := rcontest.NewServer(nil)
+	defer server.Close()
 
 	t.Run("incorrect command", func(t *testing.T) {
-		conn, err := rcon.Dial(server.Addr(), rcontest.MockPasswordRCON)
+		conn, err := rcon.Dial(server.Addr(), "password")
 		if !assert.NoError(t, err) {
 			return
 		}
@@ -71,7 +71,7 @@ func TestConn_Execute(t *testing.T) {
 	})
 
 	t.Run("closed network connection 1", func(t *testing.T) {
-		conn, err := rcon.Dial(server.Addr(), rcontest.MockPasswordRCON, rcon.SetDeadline(0))
+		conn, err := rcon.Dial(server.Addr(), "password", rcon.SetDeadline(0))
 		if !assert.NoError(t, err) {
 			return
 		}
@@ -83,7 +83,7 @@ func TestConn_Execute(t *testing.T) {
 	})
 
 	t.Run("closed network connection 2", func(t *testing.T) {
-		conn, err := rcon.Dial(server.Addr(), rcontest.MockPasswordRCON)
+		conn, err := rcon.Dial(server.Addr(), "password")
 		if !assert.NoError(t, err) {
 			return
 		}
@@ -95,7 +95,7 @@ func TestConn_Execute(t *testing.T) {
 	})
 
 	t.Run("read deadline", func(t *testing.T) {
-		conn, err := rcon.Dial(server.Addr(), rcontest.MockPasswordRCON, rcon.SetDeadline(1*time.Second))
+		conn, err := rcon.Dial(server.Addr(), "password", rcon.SetDeadline(1*time.Second))
 		if !assert.NoError(t, err) {
 			return
 		}
@@ -110,7 +110,7 @@ func TestConn_Execute(t *testing.T) {
 	})
 
 	t.Run("invalid padding", func(t *testing.T) {
-		conn, err := rcon.Dial(server.Addr(), rcontest.MockPasswordRCON)
+		conn, err := rcon.Dial(server.Addr(), "password")
 		if !assert.NoError(t, err) {
 			return
 		}
@@ -123,7 +123,7 @@ func TestConn_Execute(t *testing.T) {
 	})
 
 	t.Run("success help command", func(t *testing.T) {
-		conn, err := rcon.Dial(server.Addr(), rcontest.MockPasswordRCON)
+		conn, err := rcon.Dial(server.Addr(), "password")
 		if !assert.NoError(t, err) {
 			return
 		}
@@ -138,7 +138,7 @@ func TestConn_Execute(t *testing.T) {
 	})
 
 	t.Run("rust workaround", func(t *testing.T) {
-		conn, err := rcon.Dial(server.Addr(), rcontest.MockPasswordRCON, rcon.SetDeadline(1*time.Second))
+		conn, err := rcon.Dial(server.Addr(), "password", rcon.SetDeadline(1*time.Second))
 		if !assert.NoError(t, err) {
 			return
 		}

diff --git a/rcontest/config.go b/rcontest/config.go
@@ -0,0 +1,11 @@
+package rcontest
+
+import "time"
+
+// Config contains configuration for RCON Server.
+// TODO: Implement me or remove.
+type Config struct {
+	Password             string
+	AuthResponseDelay    time.Duration
+	CommandResponseDelay time.Duration
+}
diff --git a/rcontest/handler.go b/rcontest/handler.go
@@ -1,7 +1,8 @@
 package rcontest
 
 import (
-	"fmt"
+	"bytes"
+	"encoding/binary"
 	"io"
 	"time"
 
@@ -15,14 +16,14 @@ type Handler struct {
 }
 
 // HandlerFunc defines a function to serve RCON requests.
-type HandlerFunc func(s *Server, conn io.Writer, request *rcon.Packet) error
+type HandlerFunc func(s *Server, conn io.Writer, request *rcon.Packet)
 
-func authHandler(s *Server, conn io.Writer, request *rcon.Packet) error {
+func authHandler(s *Server, conn io.Writer, request *rcon.Packet) {
 	responseType := rcon.SERVERDATA_RESPONSE_VALUE
 	responseID := request.ID
 	responseBody := ""
 
-	if request.Body() != MockPasswordRCON {
+	if request.Body() != "password" {
 		if request.Body() == "timeout" {
 			time.Sleep(rcon.DefaultDialTimeout + 1*time.Second)
 		}
@@ -31,14 +32,42 @@ func authHandler(s *Server, conn io.Writer, request *rcon.Packet) error {
 		responseID = -1
 		responseBody = string([]byte{0x00})
 	} else {
-		_ = s.write(conn, responseID, rcon.NewPacket(responseType, responseID, responseBody))
+		_, _ = rcon.NewPacket(responseType, responseID, responseBody).WriteTo(conn)
 	}
 
 	// Auth success.
-	return s.write(conn, responseID, rcon.NewPacket(rcon.SERVERDATA_AUTH_RESPONSE, responseID, responseBody))
+	_, _ = rcon.NewPacket(rcon.SERVERDATA_AUTH_RESPONSE, responseID, responseBody).WriteTo(conn)
 }
 
-func commandHandler(s *Server, conn io.Writer, request *rcon.Packet) error {
+func commandHandler(s *Server, conn io.Writer, request *rcon.Packet) {
+	writeWithInvalidPadding := func(conn io.Writer, packet *rcon.Packet) error {
+		buffer := bytes.NewBuffer(make([]byte, 0, packet.Size+4))
+
+		if err := binary.Write(buffer, binary.LittleEndian, packet.Size); err != nil {
+			return err
+		}
+
+		if err := binary.Write(buffer, binary.LittleEndian, packet.ID); err != nil {
+			return err
+		}
+
+		if err := binary.Write(buffer, binary.LittleEndian, packet.Type); err != nil {
+			return err
+		}
+
+		// Write command body, null terminated ASCII string and an empty ASCIIZ string.
+		// Second padding byte is incorrect.
+		if _, err := buffer.Write(append([]byte(packet.Body()), 0x00, 0x01)); err != nil {
+			return err
+		}
+
+		if _, err := buffer.WriteTo(conn); err != nil {
+			return err
+		}
+
+		return nil
+	}
+
 	responseType := rcon.SERVERDATA_RESPONSE_VALUE
 	responseID := request.ID
 	responseBody := ""
@@ -52,16 +81,18 @@ func commandHandler(s *Server, conn io.Writer, request *rcon.Packet) error {
 		responseBody = request.Body()
 	case "rust":
 		// Write specific Rust package.
-		if err := s.write(conn, responseID, rcon.NewPacket(4, responseID, responseBody)); err != nil {
-			return fmt.Errorf("handle write response error: %w", err)
+		if _, err := rcon.NewPacket(4, responseID, responseBody).WriteTo(conn); err != nil {
+			return
 		}
 
 		responseBody = request.Body()
 	case "padding":
-		return s.writeWithInvalidPadding(conn, responseID, rcon.NewPacket(responseType, responseID, ""))
+		_ = writeWithInvalidPadding(conn, rcon.NewPacket(responseType, responseID, ""))
+
+		return
 	default:
 		responseBody = "unknown command"
 	}
 
-	return s.write(conn, responseID, rcon.NewPacket(responseType, responseID, responseBody))
+	_, _ = rcon.NewPacket(responseType, responseID, responseBody).WriteTo(conn)
 }