new rec added

vnc_rec/message-listeners.go

@@ -0,0 +1,77 @@
+package vnc_rec
+
+import (
+	"github.com/amitbet/vncproxy/client"
+	"github.com/amitbet/vncproxy/common"
+	"github.com/amitbet/vncproxy/logger"
+	"github.com/amitbet/vncproxy/server"
+)
+
+type ClientUpdater struct {
+	conn *client.ClientConn
+}
+
+// Consume recieves vnc-server-bound messages (Client messages) and updates the server part of the proxy
+func (cc *ClientUpdater) Consume(seg *common.RfbSegment) error {
+	logger.Tracef("ClientUpdater.Consume (vnc-server-bound): got segment type=%s bytes: %v", seg.SegmentType, seg.Bytes)
+	switch seg.SegmentType {
+
+	case common.SegmentFullyParsedClientMessage:
+		clientMsg := seg.Message.(common.ClientMessage)
+		logger.Debugf("ClientUpdater.Consume:(vnc-server-bound) got ClientMessage type=%s", clientMsg.Type())
+		switch clientMsg.Type() {
+
+		case common.SetPixelFormatMsgType:
+			// update pixel format
+			logger.Debugf("ClientUpdater.Consume: updating pixel format")
+			pixFmtMsg := clientMsg.(*server.MsgSetPixelFormat)
+			cc.conn.PixelFormat = pixFmtMsg.PF
+		}
+
+		err := clientMsg.Write(cc.conn)
+		if err != nil {
+			logger.Errorf("ClientUpdater.Consume (vnc-server-bound, SegmentFullyParsedClientMessage): problem writing to port: %s", err)
+		}
+		return err
+	}
+	return nil
+}
+
+type ServerUpdater struct {
+	conn *server.ServerConn
+}
+
+func (p *ServerUpdater) Consume(seg *common.RfbSegment) error {
+
+	logger.Debugf("WriteTo.Consume (ServerUpdater): got segment type=%s, object type:%d", seg.SegmentType, seg.UpcomingObjectType)
+	switch seg.SegmentType {
+	case common.SegmentMessageStart:
+	case common.SegmentRectSeparator:
+	case common.SegmentServerInitMessage:
+		serverInitMessage := seg.Message.(*common.ServerInit)
+		p.conn.SetHeight(serverInitMessage.FBHeight)
+		p.conn.SetWidth(serverInitMessage.FBWidth)
+		p.conn.SetDesktopName(string(serverInitMessage.NameText))
+		p.conn.SetPixelFormat(&serverInitMessage.PixelFormat)
+
+	case common.SegmentBytes:
+		logger.Debugf("WriteTo.Consume (ServerUpdater SegmentBytes): got bytes len=%d", len(seg.Bytes))
+		_, err := p.conn.Write(seg.Bytes)
+		if err != nil {
+			logger.Errorf("WriteTo.Consume (ServerUpdater SegmentBytes): problem writing to port: %s", err)
+		}
+		return err
+	case common.SegmentFullyParsedClientMessage:
+
+		clientMsg := seg.Message.(common.ClientMessage)
+		logger.Debugf("WriteTo.Consume (ServerUpdater): got ClientMessage type=%s", clientMsg.Type())
+		err := clientMsg.Write(p.conn)
+		if err != nil {
+			logger.Errorf("WriteTo.Consume (ServerUpdater SegmentFullyParsedClientMessage): problem writing to port: %s", err)
+		}
+		return err
+	default:
+		//return errors.New("WriteTo.Consume: undefined RfbSegment type")
+	}
+	return nil
+}

vnc_rec/proxy.go

@@ -0,0 +1,203 @@
+package vnc_rec
+
+import (
+	"net"
+	"path"
+	"strconv"
+	"time"
+
+	"github.com/amitbet/vncproxy/client"
+	"github.com/amitbet/vncproxy/common"
+	"github.com/amitbet/vncproxy/encodings"
+	"github.com/amitbet/vncproxy/logger"
+	"github.com/amitbet/vncproxy/player"
+	"github.com/amitbet/vncproxy/server"
+)
+
+type VncProxy struct {
+	TCPListeningURL  string      // empty = not listening on tcp
+	WsListeningURL   string      // empty = not listening on ws
+	RecordingDir     string      // empty = no recording
+	ProxyVncPassword string      //empty = no auth
+	SingleSession    *VncSession // to be used when not using sessions
+	UsingSessions    bool        //false = single session - defined in the var above
+	sessionManager   *SessionManager
+}
+
+func (vp *VncProxy) createClientConnection(target string, vncPass string) (*client.ClientConn, error) {
+	var (
+		nc  net.Conn
+		err error
+	)
+
+	if target[0] == '/' {
+		nc, err = net.Dial("unix", target)
+	} else {
+		nc, err = net.Dial("tcp", target)
+	}
+
+	if err != nil {
+		logger.Errorf("error connecting to vnc server: %s", err)
+		return nil, err
+	}
+
+	var noauth client.ClientAuthNone
+	authArr := []client.ClientAuth{&client.PasswordAuth{Password: vncPass}, &noauth}
+
+	clientConn, err := client.NewClientConn(nc,
+		&client.ClientConfig{
+			Auth:      authArr,
+			Exclusive: true,
+		})
+
+	if err != nil {
+		logger.Errorf("error creating client: %s", err)
+		return nil, err
+	}
+
+	return clientConn, nil
+}
+
+// if sessions not enabled, will always return the configured target server (only one)
+func (vp *VncProxy) getProxySession(sessionId string) (*VncSession, error) {
+
+	if !vp.UsingSessions {
+		if vp.SingleSession == nil {
+			logger.Errorf("SingleSession is empty, use sessions or populate the SingleSession member of the VncProxy struct.")
+		}
+		return vp.SingleSession, nil
+	}
+	return vp.sessionManager.GetSession(sessionId)
+}
+
+func (vp *VncProxy) newServerConnHandler(cfg *server.ServerConfig, sconn *server.ServerConn) error {
+	var err error
+	session, err := vp.getProxySession(sconn.SessionId)
+	if err != nil {
+		logger.Errorf("Proxy.newServerConnHandler can't get session: %d", sconn.SessionId)
+		return err
+	}
+
+	var rec *Recorder
+
+	if session.Type == SessionTypeRecordingProxy {
+		recFile := "recording" + strconv.FormatInt(time.Now().Unix(), 10) + ".rbs"
+		recPath := path.Join(vp.RecordingDir, recFile)
+		rec, err = NewRecorder(recPath)
+		if err != nil {
+			logger.Errorf("Proxy.newServerConnHandler can't open recorder save path: %s", recPath)
+			return err
+		}
+
+		sconn.Listeners.AddListener(rec)
+	}
+
+	session.Status = SessionStatusInit
+	if session.Type == SessionTypeProxyPass || session.Type == SessionTypeRecordingProxy {
+		target := session.Target
+		if session.TargetHostname != "" && session.TargetPort != "" {
+			target = session.TargetHostname + ":" + session.TargetPort
+		}
+
+		cconn, err := vp.createClientConnection(target, session.TargetPassword)
+		if err != nil {
+			session.Status = SessionStatusError
+			logger.Errorf("Proxy.newServerConnHandler error creating connection: %s", err)
+			return err
+		}
+		if session.Type == SessionTypeRecordingProxy {
+			cconn.Listeners.AddListener(rec)
+		}
+
+		//creating cross-listeners between server and client parts to pass messages through the proxy:
+
+		// gets the bytes from the actual vnc server on the env (client part of the proxy)
+		// and writes them through the server socket to the vnc-client
+		serverUpdater := &ServerUpdater{sconn}
+		cconn.Listeners.AddListener(serverUpdater)
+
+		// gets the messages from the server part (from vnc-client),
+		// and write through the client to the actual vnc-server
+		clientUpdater := &ClientUpdater{cconn}
+		sconn.Listeners.AddListener(clientUpdater)
+
+		err = cconn.Connect()
+		if err != nil {
+			session.Status = SessionStatusError
+			logger.Errorf("Proxy.newServerConnHandler error connecting to client: %s", err)
+			return err
+		}
+
+		encs := []common.IEncoding{
+			&encodings.RawEncoding{},
+			&encodings.TightEncoding{},
+			&encodings.EncCursorPseudo{},
+			&encodings.EncLedStatePseudo{},
+			&encodings.TightPngEncoding{},
+			&encodings.RREEncoding{},
+			&encodings.ZLibEncoding{},
+			&encodings.ZRLEEncoding{},
+			&encodings.CopyRectEncoding{},
+			&encodings.CoRREEncoding{},
+			&encodings.HextileEncoding{},
+		}
+		cconn.Encs = encs
+
+		if err != nil {
+			session.Status = SessionStatusError
+			logger.Errorf("Proxy.newServerConnHandler error connecting to client: %s", err)
+			return err
+		}
+	}
+
+	if session.Type == SessionTypeReplayServer {
+		fbs, err := player.ConnectFbsFile(session.ReplayFilePath, sconn)
+
+		if err != nil {
+			logger.Error("TestServer.NewConnHandler: Error in loading FBS: ", err)
+			return err
+		}
+		sconn.Listeners.AddListener(player.NewFBSPlayListener(sconn, fbs))
+		return nil
+
+	}
+
+	session.Status = SessionStatusActive
+	return nil
+}
+
+func (vp *VncProxy) StartListening() {
+
+	secHandlers := []server.SecurityHandler{&server.ServerAuthNone{}}
+
+	if vp.ProxyVncPassword != "" {
+		secHandlers = []server.SecurityHandler{&server.ServerAuthVNC{vp.ProxyVncPassword}}
+	}
+	cfg := &server.ServerConfig{
+		SecurityHandlers: secHandlers,
+		Encodings:        []common.IEncoding{&encodings.RawEncoding{}, &encodings.TightEncoding{}, &encodings.CopyRectEncoding{}},
+		PixelFormat:      common.NewPixelFormat(32),
+		ClientMessages:   server.DefaultClientMessages,
+		DesktopName:      []byte("workDesk"),
+		Height:           uint16(768),
+		Width:            uint16(1024),
+		NewConnHandler:   vp.newServerConnHandler,
+		UseDummySession:  !vp.UsingSessions,
+	}
+
+	if vp.TCPListeningURL != "" && vp.WsListeningURL != "" {
+		logger.Infof("running two listeners: tcp port: %s, ws url: %s", vp.TCPListeningURL, vp.WsListeningURL)
+
+		go server.WsServe(vp.WsListeningURL, cfg)
+		server.TcpServe(vp.TCPListeningURL, cfg)
+	}
+
+	if vp.WsListeningURL != "" {
+		logger.Infof("running ws listener url: %s", vp.WsListeningURL)
+		server.WsServe(vp.WsListeningURL, cfg)
+	}
+	if vp.TCPListeningURL != "" {
+		logger.Infof("running tcp listener on port: %s", vp.TCPListeningURL)
+		server.TcpServe(vp.TCPListeningURL, cfg)
+	}
+}

vnc_rec/proxy_test.go

@@ -0,0 +1,27 @@
+package vnc_rec
+
+import "testing"
+
+func TestProxy(t *testing.T) {
+	//create default session if required
+	t.Skip("this isn't an automated test, just an entrypoint for debugging")
+
+	proxy := &VncProxy{
+		WsListeningURL:  "http://0.0.0.0:7778/", // empty = not listening on ws
+		RecordingDir:    "d:\\",                 // empty = no recording
+		TCPListeningURL: ":5904",
+		//RecordingDir:          "C:\\vncRec", // empty = no recording
+		ProxyVncPassword: "1234", //empty = no auth
+		SingleSession: &VncSession{
+			TargetHostname: "192.168.1.101",
+			TargetPort:     "5901",
+			TargetPassword: "123456",
+			ID:             "dummySession",
+			Status:         SessionStatusInit,
+			Type:           SessionTypeRecordingProxy,
+		}, // to be used when not using sessions
+		UsingSessions: false, //false = single session - defined in the var above
+	}
+
+	proxy.StartListening()
+}

vnc_rec/recorder.go

@@ -0,0 +1,215 @@
+package vnc_rec
+
+import (
+	"bytes"
+	"encoding/binary"
+	"os"
+	"time"
+
+	"github.com/amitbet/vncproxy/common"
+	"github.com/amitbet/vncproxy/logger"
+	"github.com/amitbet/vncproxy/server"
+)
+
+type Recorder struct {
+	//common.BytesListener
+	RBSFileName string
+	writer      *os.File
+	//logger              common.Logger
+	startTime           int
+	buffer              bytes.Buffer
+	serverInitMessage   *common.ServerInit
+	sessionStartWritten bool
+	segmentChan         chan *common.RfbSegment
+	maxWriteSize        int
+}
+
+func getNowMillisec() int {
+	return int(time.Now().UnixNano() / int64(time.Millisecond))
+}
+
+func NewRecorder(saveFilePath string) (*Recorder, error) {
+	//delete file if it exists
+	if _, err := os.Stat(saveFilePath); err == nil {
+		os.Remove(saveFilePath)
+	}
+
+	rec := Recorder{RBSFileName: saveFilePath, startTime: getNowMillisec()}
+	var err error
+
+	rec.maxWriteSize = 65535
+
+	rec.writer, err = os.OpenFile(saveFilePath, os.O_RDWR|os.O_CREATE, 0644)
+	if err != nil {
+		logger.Errorf("unable to open file: %s, error: %v", saveFilePath, err)
+		return nil, err
+	}
+
+	//buffer the channel so we don't halt the proxying flow for slow writes when under pressure
+	rec.segmentChan = make(chan *common.RfbSegment, 100)
+	go func() {
+		for {
+			data := <-rec.segmentChan
+			rec.HandleRfbSegment(data)
+		}
+	}()
+
+	return &rec, nil
+}
+
+const versionMsg_3_3 = "RFB 003.003\n"
+const versionMsg_3_7 = "RFB 003.007\n"
+const versionMsg_3_8 = "RFB 003.008\n"
+
+// Security types
+const (
+	SecTypeInvalid = 0
+	SecTypeNone    = 1
+	SecTypeVncAuth = 2
+	SecTypeTight   = 16
+)
+
+// func (r *Recorder) writeHeader() error {
+// 	_, err := r.writer.WriteString("FBS 001.000\n")
+// 	return err
+// 	// df.write("FBS 001.000\n".getBytes());
+// }
+
+func (r *Recorder) writeStartSession(initMsg *common.ServerInit) error {
+	r.sessionStartWritten = true
+	desktopName := string(initMsg.NameText)
+	framebufferWidth := initMsg.FBWidth
+	framebufferHeight := initMsg.FBHeight
+
+	//write rfb header information (the only part done without the [size|data|timestamp] block wrapper)
+	r.writer.WriteString("FBS 001.000\n")
+
+	//push the version message into the buffer so it will be written in the first rbs block
+	r.buffer.WriteString(versionMsg_3_3)
+
+	//push sec type and fb dimensions
+	binary.Write(&r.buffer, binary.BigEndian, int32(SecTypeNone))
+	binary.Write(&r.buffer, binary.BigEndian, int16(framebufferWidth))
+	binary.Write(&r.buffer, binary.BigEndian, int16(framebufferHeight))
+
+	buff := bytes.Buffer{}
+	//binary.Write(&buff, binary.BigEndian, initMsg.FBWidth)
+	//binary.Write(&buff, binary.BigEndian, initMsg.FBHeight)
+	binary.Write(&buff, binary.BigEndian, initMsg.PixelFormat)
+	buff.Write([]byte{0, 0, 0}) //padding
+	r.buffer.Write(buff.Bytes())
+	//logger.Debugf(">>>>>>buffer for initMessage:%v ", buff.Bytes())
+
+	//var fbsServerInitMsg = []byte{32, 24, 0, 1, 0, byte(0xFF), 0, byte(0xFF), 0, byte(0xFF), 16, 8, 0, 0, 0, 0}
+	//r.buffer.Write(fbsServerInitMsg)
+
+	binary.Write(&r.buffer, binary.BigEndian, uint32(len(desktopName)))
+
+	r.buffer.WriteString(desktopName)
+	//binary.Write(&r.buffer, binary.BigEndian, byte(0)) // add null termination for desktop string
+
+	return nil
+}
+
+func (r *Recorder) Consume(data *common.RfbSegment) error {
+	//using async writes so if chan buffer overflows, proxy will not be affected
+	select {
+	case r.segmentChan <- data:
+		// default:
+		// 	logger.Error("error: recorder queue is full")
+	}
+
+	return nil
+}
+
+func (r *Recorder) HandleRfbSegment(data *common.RfbSegment) error {
+	defer func() {
+		if r := recover(); r != nil {
+			logger.Error("Recovered in HandleRfbSegment: ", r)
+		}
+	}()
+
+	switch data.SegmentType {
+	case common.SegmentMessageStart:
+		if !r.sessionStartWritten {
+			logger.Debugf("Recorder.HandleRfbSegment: writing start session segment: %v", r.serverInitMessage)
+			r.writeStartSession(r.serverInitMessage)
+		}
+
+		switch common.ServerMessageType(data.UpcomingObjectType) {
+		case common.FramebufferUpdate:
+			logger.Debugf("Recorder.HandleRfbSegment: saving FramebufferUpdate segment")
+			//r.writeToDisk()
+		case common.SetColourMapEntries:
+		case common.Bell:
+		case common.ServerCutText:
+		default:
+			logger.Warn("Recorder.HandleRfbSegment: unknown message type:" + string(data.UpcomingObjectType))
+		}
+	case common.SegmentConnectionClosed:
+		r.writeToDisk()
+	case common.SegmentRectSeparator:
+		logger.Debugf("Recorder.HandleRfbSegment: writing rect")
+		//r.writeToDisk()
+	case common.SegmentBytes:
+		logger.Debug("Recorder.HandleRfbSegment: writing bytes, len:", len(data.Bytes))
+		if r.buffer.Len()+len(data.Bytes) > r.maxWriteSize-4 {
+			r.writeToDisk()
+		}
+		_, err := r.buffer.Write(data.Bytes)
+		return err
+	case common.SegmentServerInitMessage:
+		r.serverInitMessage = data.Message.(*common.ServerInit)
+	case common.SegmentFullyParsedClientMessage:
+		clientMsg := data.Message.(common.ClientMessage)
+
+		switch clientMsg.Type() {
+		case common.SetPixelFormatMsgType:
+			clientMsg := data.Message.(*server.MsgSetPixelFormat)
+			logger.Debugf("Recorder.HandleRfbSegment: client message %v", *clientMsg)
+			r.serverInitMessage.PixelFormat = clientMsg.PF
+		default:
+			//return errors.New("unknown client message type:" + string(data.UpcomingObjectType))
+		}
+
+	default:
+		//return errors.New("undefined RfbSegment type")
+	}
+	return nil
+}
+
+func (r *Recorder) writeToDisk() error {
+	timeSinceStart := getNowMillisec() - r.startTime
+	if r.buffer.Len() == 0 {
+		return nil
+	}
+
+	//write buff length
+	bytesLen := r.buffer.Len()
+	binary.Write(r.writer, binary.BigEndian, uint32(bytesLen))
+	paddedSize := (bytesLen + 3) & 0x7FFFFFFC
+	paddingSize := paddedSize - bytesLen
+
+	//logger.Debugf("paddedSize=%d paddingSize=%d bytesLen=%d", paddedSize, paddingSize, bytesLen)
+	//write buffer padded to 32bit
+	_, err := r.buffer.WriteTo(r.writer)
+	padding := make([]byte, paddingSize)
+	//logger.Debugf("padding=%v ", padding)
+
+	binary.Write(r.writer, binary.BigEndian, padding)
+
+	//write timestamp
+	binary.Write(r.writer, binary.BigEndian, uint32(timeSinceStart))
+	r.buffer.Reset()
+	return err
+}
+
+// func (r *Recorder) WriteUInt8(data uint8) error {
+// 	buf := make([]byte, 1)
+// 	buf[0] = byte(data) // cast int8 to byte
+// 	return r.Write(buf)
+// }
+
+func (r *Recorder) Close() {
+	r.writer.Close()
+}

vnc_rec/session-manager.go

@@ -0,0 +1,19 @@
+package vnc_rec
+
+type SessionManager struct {
+	sessions map[string]*VncSession
+}
+
+func (s *SessionManager) GetSession(sessionId string) (*VncSession, error) {
+	return s.sessions[sessionId], nil
+}
+
+func (s *SessionManager) SetSession(sessionId string, session *VncSession) error {
+	s.sessions[sessionId] = session
+	return nil
+}
+
+func (s *SessionManager) DeleteSession(sessionId string) error {
+	delete(s.sessions, sessionId)
+	return nil
+}

vnc_rec/vnc-session.go

@@ -0,0 +1,27 @@
+package vnc_rec
+
+type SessionStatus int
+type SessionType int
+
+const (
+	SessionStatusInit SessionStatus = iota
+	SessionStatusActive
+	SessionStatusError
+)
+
+const (
+	SessionTypeRecordingProxy SessionType = iota
+	SessionTypeReplayServer
+	SessionTypeProxyPass
+)
+
+type VncSession struct {
+	Target         string
+	TargetHostname string
+	TargetPort     string
+	TargetPassword string
+	ID             string
+	Status         SessionStatus
+	Type           SessionType
+	ReplayFilePath string
+}