teraWeb/backend/internal/portservice/shared_gateway_test.go

package portservice

import (
	"bufio"
	"context"
	"crypto/ed25519"
	"crypto/rand"
	"fmt"
	"io"
	"net"
	"os"
	"path/filepath"
	"testing"
	"time"

	"golang.org/x/crypto/ssh"
	"tera-clone/services/tera-api/internal/directstream"
	"tera-clone/services/tera-api/internal/domain"
)

func TestBothOperationsShareCentralStreamWithDirectSession(t *testing.T) {
	for _, modes := range [][2]string{{"reverse_ssh", "reverse_telnet"}, {"reverse_ssh", "reverse_ssh"}, {"reverse_telnet", "reverse_telnet"}} {
		t.Run(modes[0]+"-"+modes[1], func(t *testing.T) {
			factory := &pipeFactory{peers: make(chan net.Conn, 1)}
			broker, _ := directstream.NewBroker(factory, directstream.BrokerOptions{SharedPortSessions: true})
			defer broker.Close()
			dir, err := os.MkdirTemp("", "gateway-")
			if err != nil {
				t.Fatal(err)
			}
			defer os.RemoveAll(dir)
			path := filepath.Join(dir, "s.sock")
			listener, err := directstream.ListenTicketSocket(path)
			if err != nil {
				t.Fatal(err)
			}
			defer listener.Close()
			go directstream.ServeSharedSocket(listener, broker)
			ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
			defer cancel()
			direct, err := broker.Attach(ctx, directstream.AttachRequest{SessionID: "web", Channel: "CH20", Actor: "web", Access: directstream.AccessWriter, SharedWrite: true, MaxAttachments: 4})
			if err != nil {
				t.Fatal(err)
			}
			defer direct.Close()
			peer := <-factory.peers
			defer peer.Close()
			_ = peer.SetDeadline(time.Now().Add(8 * time.Second))
			_, key, _ := ed25519.GenerateKey(rand.Reader)
			signer, _ := ssh.NewSignerFromKey(key)
			listeners := map[string]net.Listener{}
			gateway, err := New(Options{TransportFactory: factory, AttachShared: (directstream.SharedSocketClient{Path: path}).Attach, SSHSigner: signer,
				Authenticate: func(context.Context, string, string, string) (string, error) { return "tester", nil },
				ListenInterface: func(device, address string) (net.Listener, error) {
					if address != "0.0.0.0:10020" {
						t.Errorf("address = %s", address)
					}
					value, err := net.Listen("tcp", "127.0.0.1:0")
					listeners[device] = value
					return value, err
				},
			})
			if err != nil {
				t.Fatal(err)
			}
			defer gateway.Close()
			port := domain.Port{Channel: "CH20", Operation1: domain.PortOperation{Mode: modes[0], Port: 10020, MaxClients: 4, Authentication: "none", Newline: "cr"}, Operation2: domain.PortOperation{Mode: modes[1], Port: 10020, MaxClients: 4, Authentication: "none", Newline: "cr"}}
			if err := gateway.Apply(ctx, []domain.Port{port}); err != nil {
				t.Fatal(err)
			}
			var readers []*bufio.Reader
			for index, mode := range modes {
				address := listeners[fmt.Sprintf("eth%d", index)].Addr().String()
				var input io.Writer
				var output io.Reader
				if mode == "reverse_ssh" {
					client, err := ssh.Dial("tcp", address, &ssh.ClientConfig{User: "tester", Auth: []ssh.AuthMethod{ssh.Password("test")}, HostKeyCallback: ssh.InsecureIgnoreHostKey(), Timeout: time.Second})
					if err != nil {
						t.Fatal(err)
					}
					defer client.Close()
					session, err := client.NewSession()
					if err != nil {
						t.Fatal(err)
					}
					defer session.Close()
					input, _ = session.StdinPipe()
					output, _ = session.StdoutPipe()
					if err := session.Shell(); err != nil {
						t.Fatal(err)
					}
				} else {
					connection, err := net.DialTimeout("tcp", address, time.Second)
					if err != nil {
						t.Fatal(err)
					}
					defer connection.Close()
					_ = connection.SetDeadline(time.Now().Add(8 * time.Second))
					input, output = connection, connection
				}
				reader := bufio.NewReader(output)
				if _, err := readUntilContains(reader, "Connected at:"); err != nil {
					t.Fatal(err)
				}
				if _, err := reader.ReadString('\n'); err != nil {
					t.Fatal(err)
				}
				readers = append(readers, reader)
				want := fmt.Sprintf("op%d\r", index+1)
				go input.Write([]byte(want))
				actual := make([]byte, len(want))
				if _, err := io.ReadFull(peer, actual); err != nil || string(actual) != want {
					t.Fatalf("operation TX %q: %v", actual, err)
				}
			}
			go peer.Write([]byte("shared-output\r\n"))
			for _, reader := range readers {
				line, err := reader.ReadString('\n')
				if err != nil || line != "shared-output\r\n" {
					t.Fatalf("operation RX %q: %v", line, err)
				}
			}
			payload, err := direct.Receive(ctx)
			if err != nil || string(payload) != "shared-output\r\n" {
				t.Fatalf("web RX %q: %v", payload, err)
			}
			if snapshot := broker.Snapshot(); snapshot.TotalAttachments != 3 || snapshot.TransportOpens != 1 {
				t.Fatalf("snapshot: %+v", snapshot)
			}
		})
	}
}