added config file to workerconfig converter

slitherdmon · Jul 7, 2021 · ae5816b · ae5816b
1 parent 7b6ca20
commit ae5816b
Show file tree

Hide file tree

Showing 6 changed files with 217 additions and 77 deletions.
diff --git a/config/file.go b/config/file.go
@@ -6,6 +6,9 @@ import (
 	"log"
 	"os"
 	"path/filepath"
+
+	"github.com/realDragonium/Ultraviolet/mc"
+	"github.com/realDragonium/Ultraviolet/proxy"
 )
 
 var (
@@ -78,3 +81,18 @@ func ReadUltravioletConfig(path string) (UltravioletConfig, error) {
 	}
 	return cfg, nil
 }
+
+func FileToWorkerConfig(cfg ServerConfig) proxy.WorkerServerConfig {
+	disconPk := mc.ClientBoundDisconnect{
+		Reason: mc.Chat(cfg.DisconnectMessage),
+	}.Marshal()
+	offlineStatusPk := cfg.OfflineStatus.Marshal()
+	return proxy.WorkerServerConfig{
+		ProxyTo:           cfg.ProxyTo,
+		ProxyBind:         cfg.ProxyBind,
+		SendProxyProtocol: cfg.SendProxyProtocol,
+		ConnLimitBackend:  cfg.ConnLimitBackend,
+		OfflineStatus:     offlineStatusPk,
+		DisconnectPacket:  disconPk,
+	}
+}
diff --git a/config/file_test.go b/config/file_test.go
@@ -0,0 +1,68 @@
+package config_test
+
+import (
+	"bytes"
+	"testing"
+
+	"github.com/realDragonium/Ultraviolet/config"
+	"github.com/realDragonium/Ultraviolet/mc"
+)
+
+func samePk(expected, received mc.Packet) bool {
+	sameID := expected.ID == received.ID
+	sameData := bytes.Equal(expected.Data, received.Data)
+
+	return sameID && sameData
+}
+
+func TestFileToWorkerConfig(t *testing.T) {
+	serverCfg := config.ServerConfig{
+		MainDomain:        "Ultraviolet",
+		ExtraDomains:      []string{"Ultraviolet2", "UltraV", "UV"},
+		ProxyTo:           "127.0.10.5:25565",
+		ProxyBind:         "127.0.0.5",
+		SendProxyProtocol: true,
+		DisconnectMessage: "HelloThereWeAreClosed...Sorry",
+		OfflineStatus: mc.AnotherStatusResponse{
+			Name:        "Ultraviolet",
+			Protocol:    755,
+			Description: "Some broken proxy",
+		},
+		ConnLimitBackend: 5,
+	}
+
+	expectedDisconPk := mc.ClientBoundDisconnect{
+		Reason: mc.String(serverCfg.DisconnectMessage),
+	}.Marshal()
+	expectedOfflineStatus := mc.AnotherStatusResponse{
+		Name:        "Ultraviolet",
+		Protocol:    755,
+		Description: "Some broken proxy",
+	}.Marshal()
+
+	workerCfg := config.FileToWorkerConfig(serverCfg)
+
+	if workerCfg.ProxyTo != serverCfg.ProxyTo {
+		t.Errorf("expected: %v - got: %v", serverCfg.ProxyTo, workerCfg.ProxyTo)
+	}
+	if workerCfg.ProxyBind != serverCfg.ProxyBind {
+		t.Errorf("expected: %v - got: %v", serverCfg.ProxyBind, workerCfg.ProxyBind)
+	}
+	if workerCfg.SendProxyProtocol != serverCfg.SendProxyProtocol {
+		t.Errorf("expected: %v - got: %v", serverCfg.SendProxyProtocol, workerCfg.SendProxyProtocol)
+	}
+	if workerCfg.ConnLimitBackend != serverCfg.ConnLimitBackend {
+		t.Errorf("expected: %v - got: %v", serverCfg.ConnLimitBackend, workerCfg.ConnLimitBackend)
+	}
+	if !samePk(expectedOfflineStatus, workerCfg.OfflineStatus) {
+		offlineStatus, _ := mc.UnmarshalClientBoundResponse(expectedOfflineStatus)
+		receivedStatus, _ := mc.UnmarshalClientBoundResponse(workerCfg.OfflineStatus)
+		t.Errorf("expcted: %v \ngot: %v", offlineStatus, receivedStatus)
+	}
+
+	if !samePk(expectedDisconPk, workerCfg.DisconnectPacket) {
+		expectedDiscon, _ := mc.UnmarshalClientDisconnect(expectedDisconPk)
+		receivedDiscon, _ := mc.UnmarshalClientDisconnect(workerCfg.DisconnectPacket)
+		t.Errorf("expcted: %v \ngot: %v", expectedDiscon, receivedDiscon)
+	}
+}
diff --git a/mc/packet_login.go b/mc/packet_login.go
@@ -122,3 +122,14 @@ func (pk ClientBoundDisconnect) Marshal() Packet {
 		pk.Reason,
 	)
 }
+
+func UnmarshalClientDisconnect(packet Packet) (ClientBoundDisconnect, error) {
+	var pk ClientBoundDisconnect
+
+	if packet.ID != ClientBoundDisconnectPacketID {
+		return pk, ErrInvalidPacketID
+	}
+
+	err := packet.Scan(&pk.Reason)
+	return pk, err
+}
diff --git a/mc/packet_login_test.go b/mc/packet_login_test.go
@@ -473,3 +473,40 @@ func TestClientBoundDisconnect_Marshal(t *testing.T) {
 		}
 	}
 }
+
+func TestUnmarshalClientDisconnect(t *testing.T) {
+	tt := []struct {
+		packet             mc.Packet
+		unmarshalledPacket mc.ClientBoundDisconnect
+	}{
+		{
+			packet: mc.Packet{
+				ID:   0x00,
+				Data: []byte{0x00},
+			},
+			unmarshalledPacket: mc.ClientBoundDisconnect{
+				Reason: mc.Chat(""),
+			},
+		},
+		{
+			packet: mc.Packet{
+				ID:   0x00,
+				Data: []byte{0x0d, 0x48, 0x65, 0x6c, 0x6c, 0x6f, 0x2c, 0x20, 0x57, 0x6f, 0x72, 0x6c, 0x64, 0x21},
+			},
+			unmarshalledPacket: mc.ClientBoundDisconnect{
+				Reason: mc.Chat("Hello, World!"),
+			},
+		},
+	}
+
+	for _, tc := range tt {
+		disconnectMessage, err := mc.UnmarshalClientDisconnect(tc.packet)
+		if err != nil {
+			t.Error(err)
+		}
+
+		if disconnectMessage.Reason != tc.unmarshalledPacket.Reason {
+			t.Errorf("got: %v, want: %v", disconnectMessage.Reason, tc.unmarshalledPacket.Reason)
+		}
+	}
+}
diff --git a/proxy/worker.go b/proxy/worker.go
@@ -54,51 +54,49 @@ func SomethingElse() {
 
 }
 
-type UnknownServer struct {
-	Status mc.Packet
-}
-
 type WorkerServerConfig struct {
 	//Adding domains temporarily for testing until better structure
 	MainDomain   string
 	ExtraDomains []string
 
-	OnlineStatus  mc.Packet
-	OfflineStatus mc.Packet
-	State         ServerState
+	State ServerState
+
+	OnlineStatus     mc.Packet
+	OfflineStatus    mc.Packet
+	DisconnectPacket mc.Packet
 
 	ProxyTo           string
 	ProxyBind         string
 	SendProxyProtocol bool
-	DisconnectMessage string
 
 	ConnLimitBackend int
 }
 
-func NewWorker(req chan conn.ConnRequest) Worker {
+func NewWorker(req chan conn.ConnRequest, proxies map[string]WorkerServerConfig, defaultStatus mc.Packet) Worker {
 	return Worker{
-		ReqCh: req,
+		reqCh:         req,
+		defaultStatus: defaultStatus,
+		Servers:       proxies,
 	}
-
 }
 
 type Worker struct {
-	ReqCh      chan conn.ConnRequest
-	DefaultCfg UnknownServer
-	Servers    map[string]WorkerServerConfig
+	reqCh         chan conn.ConnRequest
+	defaultStatus mc.Packet
+	Servers       map[string]WorkerServerConfig
 }
 
 func (w Worker) Work() {
 	for {
-		request := <-w.ReqCh
+		request := <-w.reqCh
 		serverCfg, ok := w.Servers[request.ServerAddr]
 		if !ok {
 			//Unknown server address
 			switch request.Type {
 			case conn.STATUS:
 				request.Ch <- conn.ConnAnswer{
 					Action:   conn.SEND_STATUS,
-					StatusPk: w.DefaultCfg.Status,
+					StatusPk: w.defaultStatus,
 				}
 			case conn.LOGIN:
 				request.Ch <- conn.ConnAnswer{
@@ -120,7 +118,8 @@ func (w Worker) Work() {
 		case conn.LOGIN:
 			if serverCfg.State == OFFLINE {
 				request.Ch <- conn.ConnAnswer{
-					Action: conn.CLOSE,
+					Action:        conn.DISCONNECT,
+					DisconMessage: serverCfg.DisconnectPacket,
 				}
 				return
 			}