Skip to content
Snippets Groups Projects
protocol_test.go 5.79 KiB
Newer Older
  • Learn to ignore specific revisions
  • Bolke de Bruin's avatar
    Bolke de Bruin committed
    package protocol
    
    import (
    
    Bolke de Bruin's avatar
    Bolke de Bruin committed
    	"log"
    	"testing"
    )
    
    const (
    
    Bolke de Bruin's avatar
    Bolke de Bruin committed
    	HeaderLen               = 8
    	HandshakeRequestLen     = HeaderLen + 6
    	HandshakeResponseLen    = HeaderLen + 10
    	TunnelCreateRequestLen  = HeaderLen + 8 // + dynamic
    
    	TunnelCreateResponseLen = HeaderLen + 18
    
    Bolke de Bruin's avatar
    Bolke de Bruin committed
    	TunnelAuthLen           = HeaderLen + 2 // + dynamic
    	TunnelAuthResponseLen   = HeaderLen + 16
    
    Bolke de Bruin's avatar
    Bolke de Bruin committed
    	ChannelCreateLen		= HeaderLen + 8 // + dynamic
    	ChannelResponseLen		= HeaderLen + 12
    
    Bolke de Bruin's avatar
    Bolke de Bruin committed
    )
    
    
    Bolke de Bruin's avatar
    Bolke de Bruin committed
    func verifyPacketHeader(data []byte, expPt uint16, expSize uint32) (uint16, uint32, []byte, error) {
    
    	pt, size, pkt, err := readHeader(data)
    
    	if pt != expPt {
    
    Bolke de Bruin's avatar
    Bolke de Bruin committed
    		return 0, 0, []byte{}, fmt.Errorf("readHeader failed, expected packet type %d got %d", expPt, pt)
    
    	}
    
    	if size != expSize {
    		return 0, 0, []byte{}, fmt.Errorf("readHeader failed, expected size %d, got %d", expSize, size)
    	}
    
    	if err != nil {
    		return 0, 0, []byte{}, err
    	}
    
    	return pt, size, pkt, nil
    }
    
    
    Bolke de Bruin's avatar
    Bolke de Bruin committed
    func TestHandshake(t *testing.T) {
    	client := ClientConfig{
    		PAAToken: "abab",
    	}
    
    Bolke de Bruin's avatar
    Bolke de Bruin committed
    	s := &SessionInfo{}
    	hc := &ServerConf{
    		TokenAuth: true,
    	}
    	h := NewServer(s, hc)
    
    Bolke de Bruin's avatar
    Bolke de Bruin committed
    
    	data := client.handshakeRequest()
    
    
    	_, _, pkt, err := verifyPacketHeader(data, PKT_TYPE_HANDSHAKE_REQUEST, HandshakeRequestLen)
    
    Bolke de Bruin's avatar
    Bolke de Bruin committed
    
    	if err != nil {
    
    		t.Fatalf("verifyHeader failed: %s", err)
    
    Bolke de Bruin's avatar
    Bolke de Bruin committed
    	}
    
    	log.Printf("pkt: %x", pkt)
    
    
    Bolke de Bruin's avatar
    Bolke de Bruin committed
    	major, minor, version, extAuth := h.handshakeRequest(pkt)
    
    Bolke de Bruin's avatar
    Bolke de Bruin committed
    	if major != MajorVersion || minor != MinorVersion || version != Version {
    
    Bolke de Bruin's avatar
    Bolke de Bruin committed
    		t.Fatalf("handshakeRequest failed got version %d.%d protocol %d, expected %d.%d protocol %d",
    
    Bolke de Bruin's avatar
    Bolke de Bruin committed
    			major, minor, version, MajorVersion, MinorVersion, Version)
    	}
    
    	if !((extAuth & HTTP_EXTENDED_AUTH_PAA) == HTTP_EXTENDED_AUTH_PAA) {
    
    Bolke de Bruin's avatar
    Bolke de Bruin committed
    		t.Fatalf("handshakeRequest failed got ext auth %d, expected %d", extAuth, extAuth|HTTP_EXTENDED_AUTH_PAA)
    
    	}
    
    	data = h.handshakeResponse(0x0, 0x0)
    	_, _, pkt, err = verifyPacketHeader(data, PKT_TYPE_HANDSHAKE_RESPONSE, HandshakeResponseLen)
    	if err != nil {
    		t.Fatalf("verifyHeader failed: %s", err)
    	}
    	log.Printf("pkt: %x", pkt)
    
    	caps, err := client.handshakeResponse(pkt)
    	if !((caps & HTTP_EXTENDED_AUTH_PAA) == HTTP_EXTENDED_AUTH_PAA) {
    
    Bolke de Bruin's avatar
    Bolke de Bruin committed
    		t.Fatalf("handshakeResponse failed got caps %d, expected %d", caps, caps|HTTP_EXTENDED_AUTH_PAA)
    
    Bolke de Bruin's avatar
    Bolke de Bruin committed
    }
    
    
    func TestTunnelCreation(t *testing.T) {
    	client := ClientConfig{
    		PAAToken: "abab",
    	}
    
    Bolke de Bruin's avatar
    Bolke de Bruin committed
    	s := &SessionInfo{}
    	hc := &ServerConf{
    		TokenAuth: true,
    	}
    	h := NewServer(s, hc)
    
    
    	data := client.tunnelRequest()
    	_, _, pkt, err := verifyPacketHeader(data, PKT_TYPE_TUNNEL_CREATE,
    
    Bolke de Bruin's avatar
    Bolke de Bruin committed
    		uint32(TunnelCreateRequestLen+2+len(client.PAAToken)*2))
    
    	if err != nil {
    		t.Fatalf("verifyHeader failed: %s", err)
    	}
    
    
    Bolke de Bruin's avatar
    Bolke de Bruin committed
    	caps, token := h.tunnelRequest(pkt)
    
    	if !((caps & HTTP_CAPABILITY_IDLE_TIMEOUT) == HTTP_CAPABILITY_IDLE_TIMEOUT) {
    
    Bolke de Bruin's avatar
    Bolke de Bruin committed
    		t.Fatalf("tunnelRequest failed got caps %d, expected %d", caps, caps|HTTP_CAPABILITY_IDLE_TIMEOUT)
    
    	}
    	if token != client.PAAToken {
    
    Bolke de Bruin's avatar
    Bolke de Bruin committed
    		t.Fatalf("tunnelRequest failed got token %s, expected %s", token, client.PAAToken)
    
    Bolke de Bruin's avatar
    Bolke de Bruin committed
    	data = h.tunnelResponse()
    
    	_, _, pkt, err = verifyPacketHeader(data, PKT_TYPE_TUNNEL_RESPONSE, TunnelCreateResponseLen)
    	if err != nil {
    		t.Fatalf("verifyHeader failed: %s", err)
    	}
    
    	tid, caps, err := client.tunnelResponse(pkt)
    	if err != nil {
    		t.Fatalf("Error %s", err)
    	}
    	if tid != tunnelId {
    		t.Fatalf("tunnelResponse failed tunnel id %d, expected %d", tid, tunnelId)
    	}
    	if !((caps & HTTP_CAPABILITY_IDLE_TIMEOUT) == HTTP_CAPABILITY_IDLE_TIMEOUT) {
    
    Bolke de Bruin's avatar
    Bolke de Bruin committed
    		t.Fatalf("tunnelResponse failed got caps %d, expected %d", caps, caps|HTTP_CAPABILITY_IDLE_TIMEOUT)
    
    Bolke de Bruin's avatar
    Bolke de Bruin committed
    }
    
    func TestTunnelAuth(t *testing.T) {
    	client := ClientConfig{}
    	s := &SessionInfo{}
    
    Bolke de Bruin's avatar
    Bolke de Bruin committed
    	hc := &ServerConf{
    		TokenAuth:   true,
    
    Bolke de Bruin's avatar
    Bolke de Bruin committed
    		IdleTimeout: 10,
    		RedirectFlags: RedirectFlags{
    			Clipboard: true,
    		},
    	}
    
    Bolke de Bruin's avatar
    Bolke de Bruin committed
    	h := NewServer(s, hc)
    
    Bolke de Bruin's avatar
    Bolke de Bruin committed
    	name := "test_name"
    
    	data := client.tunnelAuthRequest(name)
    
    Bolke de Bruin's avatar
    Bolke de Bruin committed
    	_, _, pkt, err := verifyPacketHeader(data, PKT_TYPE_TUNNEL_AUTH, uint32(TunnelAuthLen+len(name)*2))
    
    Bolke de Bruin's avatar
    Bolke de Bruin committed
    	if err != nil {
    		t.Fatalf("verifyHeader failed: %s", err)
    	}
    
    
    Bolke de Bruin's avatar
    Bolke de Bruin committed
    	n := h.tunnelAuthRequest(pkt)
    
    Bolke de Bruin's avatar
    Bolke de Bruin committed
    	if n != name {
    
    Bolke de Bruin's avatar
    Bolke de Bruin committed
    		t.Fatalf("tunnelAuthRequest failed got name %s, expected %s", n, name)
    
    Bolke de Bruin's avatar
    Bolke de Bruin committed
    	data = h.tunnelAuthResponse()
    
    Bolke de Bruin's avatar
    Bolke de Bruin committed
    	_, _, pkt, err = verifyPacketHeader(data, PKT_TYPE_TUNNEL_AUTH_RESPONSE, TunnelAuthResponseLen)
    	if err != nil {
    		t.Fatalf("verifyHeader failed: %s", err)
    	}
    	flags, timeout, err := client.tunnelAuthResponse(pkt)
    	if err != nil {
    		t.Fatalf("tunnel auth error %s", err)
    	}
    	if (flags & HTTP_TUNNEL_REDIR_DISABLE_CLIPBOARD) == HTTP_TUNNEL_REDIR_DISABLE_CLIPBOARD {
    		t.Fatalf("tunnelAuthResponse failed got flags %d, expected %d",
    
    Bolke de Bruin's avatar
    Bolke de Bruin committed
    			flags, flags|HTTP_TUNNEL_REDIR_DISABLE_CLIPBOARD)
    
    Bolke de Bruin's avatar
    Bolke de Bruin committed
    	}
    	if int(timeout) != hc.IdleTimeout {
    		t.Fatalf("tunnelAuthResponse failed got timeout %d, expected %d",
    			timeout, hc.IdleTimeout)
    	}
    
    Bolke de Bruin's avatar
    Bolke de Bruin committed
    }
    
    Bolke de Bruin's avatar
    Bolke de Bruin committed
    
    func TestChannelCreation(t *testing.T) {
    	client := ClientConfig{}
    	s := &SessionInfo{}
    	hc := &ServerConf{
    		TokenAuth:   true,
    		IdleTimeout: 10,
    		RedirectFlags: RedirectFlags{
    			Clipboard: true,
    		},
    	}
    	h := NewServer(s, hc)
    	server := "test_server"
    	port := uint16(3389)
    
    	data := client.channelRequest(server, port)
    	_, _, pkt, err := verifyPacketHeader(data, PKT_TYPE_CHANNEL_CREATE, uint32(ChannelCreateLen+len(server)*2))
    	if err != nil {
    		t.Fatalf("verifyHeader failed: %s", err)
    	}
    	hServer, hPort := h.channelRequest(pkt)
    	if hServer != server {
    		t.Fatalf("channelRequest failed got server %s, expected %s", hServer, server)
    	}
    	if hPort != port {
    		t.Fatalf("channelRequest failed got port %d, expected %d", hPort, port)
    	}
    
    	data = h.channelResponse()
    	_, _, pkt, err = verifyPacketHeader(data, PKT_TYPE_CHANNEL_RESPONSE, uint32(ChannelResponseLen))
    	if err != nil {
    		t.Fatalf("verifyHeader failed: %s", err)
    	}
    	channelId, err := client.channelResponse(pkt)
    	if err != nil {
    		t.Fatalf("channelResponse failed: %s", err)
    	}
    	if channelId < 1 {
    		t.Fatalf("channelResponse failed got channeld id %d, expected > 0", channelId)
    	}
    
    Bolke de Bruin's avatar
    Bolke de Bruin committed
    }