diff --git a/rdg.go b/rdg.go
index 40b018a5da1f226651234a73b9093f8bf04522de..94326450784eff21b6802aeb8c12b64b2b1a291c 100644
--- a/rdg.go
+++ b/rdg.go
@@ -6,18 +6,18 @@ import (
"encoding/binary"
"errors"
"fmt"
+ "github.com/gorilla/websocket"
"github.com/patrickmn/go-cache"
"io"
"log"
"math/rand"
"net"
"net/http"
- //"net/http/httputil"
"strconv"
+
"time"
"unicode/utf16"
"unicode/utf8"
- "github.com/gorilla/websocket"
)
const (
@@ -110,7 +110,7 @@ var ErrNotHijacker = RejectConnectionError(
var DefaultSession RdgSession
func Upgrade(next http.Handler) http.Handler {
- return RdgHandshake(next)
+ return handleGatewayProtocol(next)
}
func Accept(w http.ResponseWriter) (conn net.Conn, rw *bufio.ReadWriter, err error) {
@@ -129,10 +129,9 @@ func Accept(w http.ResponseWriter) (conn net.Conn, rw *bufio.ReadWriter, err err
}
var upgrader = websocket.Upgrader{}
+var c = cache.New(5*time.Minute, 10*time.Minute)
-func RdgHandshake(next http.Handler) http.Handler {
- c := cache.New(5*time.Minute, 10*time.Minute)
-
+func handleGatewayProtocol(next http.Handler) http.Handler {
return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
var s RdgSession
@@ -150,88 +149,21 @@ func RdgHandshake(next http.Handler) http.Handler {
if r.Method == MethodRDGOUT {
r.Method = "GET" // force
- c, err := upgrader.Upgrade(w, r, nil)
+ conn, err := upgrader.Upgrade(w, r, nil)
if err != nil {
log.Printf("Cannot upgrade falling back to old protocol: %s", err)
return
}
- defer c.Close()
-
- fragment := false
- buf := make([]byte, 4096)
- index := 0
- for {
- mt, msg, err := c.ReadMessage()
- if err != nil {
- log.Printf("Error read: %s", err)
- break
- }
- log.Printf("Message type: %d, message: %x", mt, msg)
+ defer conn.Close()
- // check for fragments
- var pt uint16
- var sz uint32
- var pkt []byte
+ handleWebsocketProtocol(conn)
- if !fragment {
- pt, sz, pkt, err = readHeader(msg)
- if err != nil {
- // fragment received
- log.Printf("Received non websocket fragment")
- fragment = true
- index = copy(buf, msg)
- continue
- }
- index = 0
- } else {
- log.Printf("Dealing with fragment")
- fragment = false
- pt, sz, pkt, _ = readHeader(append(buf[:index], msg...))
- }
- switch pt {
- case PKT_TYPE_HANDSHAKE_REQUEST:
- major, minor, _, auth := readHandshake(pkt)
- msg := handshakeResponse(major, minor, auth)
- log.Printf("Handshake response: %x", msg)
- c.WriteMessage(mt, msg)
- case PKT_TYPE_TUNNEL_CREATE:
- readCreateTunnelRequest(pkt)
- msg := createTunnelResponse()
- log.Printf("Create tunnel response: %x", msg)
- c.WriteMessage(mt, msg)
- case PKT_TYPE_TUNNEL_AUTH:
- readTunnelAuthRequest(pkt)
- msg := createTunnelAuthResponse()
- log.Printf("Create tunnel auth response: %x", msg)
- c.WriteMessage(mt, msg)
- case PKT_TYPE_CHANNEL_CREATE:
- server, port := readChannelCreateRequest(pkt)
- s.Remote, err = net.Dial("tcp", net.JoinHostPort(server, strconv.Itoa(int(port))))
- if err != nil {
- log.Printf("Error connecting to %s, %d, %s", server, port, err)
- return
- }
- msg := createChannelCreateResponse()
- log.Printf("Create channel create response: %x", msg)
- c.WriteMessage(mt, msg)
- go handleWebsocketData(s.Remote, mt, c)
- case PKT_TYPE_DATA:
- forwardDataPacket(s.Remote, pkt)
- case PKT_TYPE_KEEPALIVE:
- c.WriteMessage(mt, createPacket(PKT_TYPE_KEEPALIVE, []byte{}))
- case PKT_TYPE_CLOSE_CHANNEL:
- s.Remote.Close()
- return
- default:
- log.Printf("Unknown packet type: %d (size: %d), %x", pt, sz)
- }
- }
- conn, rw, _ := Accept(w)
- log.Printf("Opening RDGOUT for client %s", conn.RemoteAddr().String())
+ //conn, rw, _ := Accept(w)
+ //log.Printf("Opening RDGOUT for client %s", conn.RemoteAddr().String())
- s.ConnOut = conn
- WriteAcceptSeed(rw.Writer, true)
+ //s.ConnOut = conn
+ //WriteAcceptSeed(rw.Writer, true)
//c.Set(connId, s, cache.DefaultExpiration)
} /*else if r.Method == MethodRDGIN {
@@ -303,6 +235,82 @@ func RdgHandshake(next http.Handler) http.Handler {
})
}
+func handleWebsocketProtocol(conn *websocket.Conn) {
+ fragment := false
+ buf := make([]byte, 4096)
+ index := 0
+
+ var remote net.Conn
+
+ for {
+ mt, msg, err := conn.ReadMessage()
+ if err != nil {
+ log.Printf("Error read: %s", err)
+ break
+ }
+ log.Printf("Message type: %d, message: %x", mt, msg)
+
+ // check for fragments
+ var pt uint16
+ var sz uint32
+ var pkt []byte
+
+ if !fragment {
+ pt, sz, pkt, err = readHeader(msg)
+ if err != nil {
+ // fragment received
+ log.Printf("Received non websocket fragment")
+ fragment = true
+ index = copy(buf, msg)
+ continue
+ }
+ index = 0
+ } else {
+ log.Printf("Dealing with fragment")
+ fragment = false
+ pt, sz, pkt, _ = readHeader(append(buf[:index], msg...))
+ }
+
+ switch pt {
+ case PKT_TYPE_HANDSHAKE_REQUEST:
+ major, minor, _, auth := readHandshake(pkt)
+ msg := handshakeResponse(major, minor, auth)
+ log.Printf("Handshake response: %x", msg)
+ conn.WriteMessage(mt, msg)
+ case PKT_TYPE_TUNNEL_CREATE:
+ readCreateTunnelRequest(pkt)
+ msg := createTunnelResponse()
+ log.Printf("Create tunnel response: %x", msg)
+ conn.WriteMessage(mt, msg)
+ case PKT_TYPE_TUNNEL_AUTH:
+ readTunnelAuthRequest(pkt)
+ msg := createTunnelAuthResponse()
+ log.Printf("Create tunnel auth response: %x", msg)
+ conn.WriteMessage(mt, msg)
+ case PKT_TYPE_CHANNEL_CREATE:
+ server, port := readChannelCreateRequest(pkt)
+ remote, err = net.Dial("tcp", net.JoinHostPort(server, strconv.Itoa(int(port))))
+ if err != nil {
+ log.Printf("Error connecting to %s, %d, %s", server, port, err)
+ return
+ }
+ msg := createChannelCreateResponse()
+ log.Printf("Create channel create response: %x", msg)
+ conn.WriteMessage(mt, msg)
+ go handleWebsocketData(remote, mt, conn)
+ case PKT_TYPE_DATA:
+ forwardDataPacket(remote, pkt)
+ case PKT_TYPE_KEEPALIVE:
+ conn.WriteMessage(mt, createPacket(PKT_TYPE_KEEPALIVE, []byte{}))
+ case PKT_TYPE_CLOSE_CHANNEL:
+ remote.Close()
+ return
+ default:
+ log.Printf("Unknown packet type: %d (size: %d), %x", pt, sz)
+ }
+ }
+}
+
// [MS-TSGU]: Terminal Services Gateway Server Protocol version 39.0
// The server sends back the final status code 200 OK, and also a random entity body of limited size (100 bytes).
// This enables a reverse proxy to start allowing data from the RDG server to the RDG client. The RDG server does