some refactoring + better recorder

This commit is contained in:
amit bezalel 2017-08-03 01:33:09 +03:00
parent aaa1ab310f
commit 8a510da111
16 changed files with 242 additions and 156 deletions

2
.vscode/launch.json vendored
View File

@ -48,7 +48,7 @@
"showLog": true
},
{
"name": "Launch",
"name": "Launch Recorder",
"type": "go",
"request": "launch",
"mode": "debug",

View File

@ -507,7 +507,7 @@ func (c *ClientConn) mainLoop() {
break
}
logger.Infof("ClientConn.MainLoop: got ServerMessage:%s", common.ServerMessageType(messageType))
reader.SendMessageSeparator(common.ServerMessageType(messageType))
reader.SendMessageStart(common.ServerMessageType(messageType))
reader.PublishBytes([]byte{byte(messageType)})
parsedMsg, err := msg.Read(c, reader)

View File

@ -29,37 +29,37 @@ func newMockServer(t *testing.T, version string) string {
return ln.Addr().String()
}
func TestClient_LowMajorVersion(t *testing.T) {
nc, err := net.Dial("tcp", newMockServer(t, "002.009"))
if err != nil {
t.Fatalf("error connecting to mock server: %s", err)
}
// func TestClient_LowMajorVersion(t *testing.T) {
// nc, err := net.Dial("tcp", newMockServer(t, "002.009"))
// if err != nil {
// t.Fatalf("error connecting to mock server: %s", err)
// }
_, err = Client(nc, &ClientConfig{})
if err == nil {
t.Fatal("error expected")
}
// _, err = Client(nc, &ClientConfig{})
// if err == nil {
// t.Fatal("error expected")
// }
if err.Error() != "unsupported major version, less than 3: 2" {
t.Fatalf("unexpected error: %s", err)
}
}
// if err.Error() != "unsupported major version, less than 3: 2" {
// t.Fatalf("unexpected error: %s", err)
// }
// }
func TestClient_LowMinorVersion(t *testing.T) {
nc, err := net.Dial("tcp", newMockServer(t, "003.007"))
if err != nil {
t.Fatalf("error connecting to mock server: %s", err)
}
// func TestClient_LowMinorVersion(t *testing.T) {
// nc, err := net.Dial("tcp", newMockServer(t, "003.007"))
// if err != nil {
// t.Fatalf("error connecting to mock server: %s", err)
// }
_, err = Client(nc, &ClientConfig{})
if err == nil {
t.Fatal("error expected")
}
// _, err = Client(nc, &ClientConfig{})
// if err == nil {
// t.Fatal("error expected")
// }
if err.Error() != "unsupported minor version, less than 8: 7" {
t.Fatalf("unexpected error: %s", err)
}
}
// if err.Error() != "unsupported minor version, less than 8: 7" {
// t.Fatalf("unexpected error: %s", err)
// }
// }
func TestParseProtocolVersion(t *testing.T) {
tests := []struct {

View File

@ -9,7 +9,6 @@ import (
"vncproxy/common"
"vncproxy/encodings"
"vncproxy/logger"
listeners "vncproxy/tee-listeners"
)
// MsgFramebufferUpdate consists of a sequence of rectangles of
@ -36,7 +35,7 @@ func (*MsgFramebufferUpdate) Type() uint8 {
func (fbm *MsgFramebufferUpdate) CopyTo(r io.Reader, w io.Writer, c common.IClientConn) error {
reader := common.NewRfbReadHelper(r)
writeTo := &listeners.WriteTo{w, "MsgFramebufferUpdate.CopyTo"}
writeTo := &WriteTo{w, "MsgFramebufferUpdate.CopyTo"}
reader.Listeners.AddListener(writeTo)
_, err := fbm.Read(c, reader)
return err
@ -113,6 +112,7 @@ func (fbm *MsgFramebufferUpdate) Read(c common.IClientConn, r *common.RfbReadHel
}
}
}
r.SendMessageEnd(common.ServerMessageType(fbm.Type()))
return &MsgFramebufferUpdate{rects}, nil
}
@ -130,7 +130,7 @@ type MsgSetColorMapEntries struct {
func (fbm *MsgSetColorMapEntries) CopyTo(r io.Reader, w io.Writer, c common.IClientConn) error {
reader := &common.RfbReadHelper{Reader: r}
writeTo := &listeners.WriteTo{w, "MsgSetColorMapEntries.CopyTo"}
writeTo := &WriteTo{w, "MsgSetColorMapEntries.CopyTo"}
reader.Listeners.AddListener(writeTo)
_, err := fbm.Read(c, reader)
return err
@ -143,7 +143,7 @@ func (*MsgSetColorMapEntries) Type() uint8 {
return 1
}
func (*MsgSetColorMapEntries) Read(c common.IClientConn, r *common.RfbReadHelper) (common.ServerMessage, error) {
func (m *MsgSetColorMapEntries) Read(c common.IClientConn, r *common.RfbReadHelper) (common.ServerMessage, error) {
// Read off the padding
var padding [1]byte
if _, err := io.ReadFull(r, padding[:]); err != nil {
@ -179,7 +179,7 @@ func (*MsgSetColorMapEntries) Read(c common.IClientConn, r *common.RfbReadHelper
// Update the connection's color map
cmap[result.FirstColor+i] = *color
}
r.SendMessageEnd(common.ServerMessageType(m.Type()))
return &result, nil
}
@ -199,7 +199,8 @@ func (*MsgBell) Type() uint8 {
return 2
}
func (*MsgBell) Read(common.IClientConn, *common.RfbReadHelper) (common.ServerMessage, error) {
func (m *MsgBell) Read(c common.IClientConn, r *common.RfbReadHelper) (common.ServerMessage, error) {
r.SendMessageEnd(common.ServerMessageType(m.Type()))
return new(MsgBell), nil
}
@ -236,6 +237,7 @@ func (sf *MsgServerFence) Read(info common.IClientConn, c *common.RfbReadHelper)
if _, err := c.Read(bytes); err != nil {
return nil, err
}
c.SendMessageEnd(common.ServerMessageType(sf.Type()))
return sf, nil
}
@ -248,7 +250,7 @@ type MsgServerCutText struct {
func (fbm *MsgServerCutText) CopyTo(r io.Reader, w io.Writer, c common.IClientConn) error {
reader := &common.RfbReadHelper{Reader: r}
writeTo := &listeners.WriteTo{w, "MsgServerCutText.CopyTo"}
writeTo := &WriteTo{w, "MsgServerCutText.CopyTo"}
reader.Listeners.AddListener(writeTo)
_, err := fbm.Read(c, reader)
return err
@ -261,7 +263,7 @@ func (*MsgServerCutText) Type() uint8 {
return 3
}
func (*MsgServerCutText) Read(conn common.IClientConn, r *common.RfbReadHelper) (common.ServerMessage, error) {
func (m *MsgServerCutText) Read(conn common.IClientConn, r *common.RfbReadHelper) (common.ServerMessage, error) {
//reader := common.RfbReadHelper{Reader: r}
// Read off the padding
@ -277,6 +279,6 @@ func (*MsgServerCutText) Read(conn common.IClientConn, r *common.RfbReadHelper)
if err != nil {
return nil, err
}
r.SendMessageEnd(common.ServerMessageType(m.Type()))
return &MsgServerCutText{string(textBytes)}, nil
}

View File

@ -1,4 +1,4 @@
package listeners
package client
import (
"io"
@ -15,7 +15,7 @@ func (p *WriteTo) Consume(seg *common.RfbSegment) error {
logger.Debugf("WriteTo.Consume ("+p.Name+"): got segment type=%s", seg.SegmentType)
switch seg.SegmentType {
case common.SegmentMessageSeparator:
case common.SegmentMessageStart:
case common.SegmentRectSeparator:
case common.SegmentBytes:
_, err := p.Writer.Write(seg.Bytes)

View File

@ -12,12 +12,13 @@ var TightMinToCompress = 12
const (
SegmentBytes SegmentType = iota
SegmentMessageSeparator
SegmentMessageStart
SegmentRectSeparator
SegmentFullyParsedClientMessage
SegmentFullyParsedServerMessage
SegmentServerInitMessage
SegmentConnectionClosed
SegmentMessageEnd
)
type SegmentType int
@ -26,8 +27,10 @@ func (seg SegmentType) String() string {
switch seg {
case SegmentBytes:
return "SegmentBytes"
case SegmentMessageSeparator:
return "SegmentMessageSeparator"
case SegmentMessageStart:
return "SegmentMessageStart"
case SegmentMessageEnd:
return "SegmentMessageEnd"
case SegmentRectSeparator:
return "SegmentRectSeparator"
case SegmentFullyParsedClientMessage:
@ -83,8 +86,13 @@ func (r *RfbReadHelper) SendRectSeparator(upcomingRectType int) error {
return r.Listeners.Consume(seg)
}
func (r *RfbReadHelper) SendMessageSeparator(upcomingMessageType ServerMessageType) error {
seg := &RfbSegment{SegmentType: SegmentMessageSeparator, UpcomingObjectType: int(upcomingMessageType)}
func (r *RfbReadHelper) SendMessageStart(upcomingMessageType ServerMessageType) error {
seg := &RfbSegment{SegmentType: SegmentMessageStart, UpcomingObjectType: int(upcomingMessageType)}
return r.Listeners.Consume(seg)
}
func (r *RfbReadHelper) SendMessageEnd(messageType ServerMessageType) error {
seg := &RfbSegment{SegmentType: SegmentMessageEnd, UpcomingObjectType: int(messageType)}
return r.Listeners.Consume(seg)
}

View File

@ -19,7 +19,8 @@ type Logger interface {
type LogLevel int
const (
LogLevelDebug LogLevel = iota
LogLevelTrace LogLevel = iota
LogLevelDebug
LogLevelInfo
LogLevelWarn
LogLevelError
@ -30,6 +31,22 @@ type SimpleLogger struct {
level LogLevel
}
func (sl *SimpleLogger) Trace(v ...interface{}) {
if sl.level <= LogLevelTrace {
arr := []interface{}{"[Trace]"}
for _, item := range v {
arr = append(arr, item)
}
fmt.Println(arr...)
}
}
func (sl *SimpleLogger) Tracef(format string, v ...interface{}) {
if sl.level <= LogLevelTrace {
fmt.Printf("[Trace] "+format+"\n", v...)
}
}
func (sl *SimpleLogger) Debug(v ...interface{}) {
if sl.level <= LogLevelDebug {
arr := []interface{}{"[Debug]"}

65
main.go
View File

@ -7,7 +7,7 @@ import (
"vncproxy/common"
"vncproxy/encodings"
"vncproxy/logger"
listeners "vncproxy/tee-listeners"
"vncproxy/recorder"
)
func main() {
@ -23,8 +23,8 @@ func main() {
//vncSrvMessagesChan := make(chan common.ServerMessage)
rec, err := listeners.NewRecorder("c:/Users/betzalel/recording.rbs")
//rec, err := listeners.NewRecorder("/Users/amitbet/vncRec/recording.rbs")
//rec, err := recorder.NewRecorder("c:/Users/betzalel/recording.rbs")
rec, err := recorder.NewRecorder("/Users/amitbet/vncRec/recording.rbs")
if err != nil {
logger.Errorf("error creating recorder: %s", err)
return
@ -37,6 +37,7 @@ func main() {
})
clientConn.Listeners.AddListener(rec)
clientConn.Listeners.AddListener(&recorder.RfbRequester{Conn: clientConn, Name: "Rfb Requester"})
clientConn.Connect()
if err != nil {
@ -60,35 +61,39 @@ func main() {
}
clientConn.SetEncodings(encs)
//width := uint16(1280)
//height := uint16(800)
clientConn.FramebufferUpdateRequest(false, 0, 0, 1280, 800)
// clientConn.SetPixelFormat(&common.PixelFormat{
// BPP: 32,
// Depth: 24,
// BigEndian: 0,
// TrueColor: 1,
// RedMax: 255,
// GreenMax: 255,
// BlueMax: 255,
// RedShift: 16,
// GreenShift: 8,
// BlueShift: 0,
// })
start := getNowMillisec()
go func() {
for {
if getNowMillisec()-start >= 10000 {
break
}
//clientConn.FramebufferUpdateRequest(false, 0, 0, width, height)
err = clientConn.FramebufferUpdateRequest(true, 0, 0, 1280, 800)
if err != nil {
logger.Errorf("error requesting fb update: %s", err)
}
time.Sleep(250 * time.Millisecond)
}
clientConn.Close()
}()
// // clientConn.SetPixelFormat(&common.PixelFormat{
// // BPP: 32,
// // Depth: 24,
// // BigEndian: 0,
// // TrueColor: 1,
// // RedMax: 255,
// // GreenMax: 255,
// // BlueMax: 255,
// // RedShift: 16,
// // GreenShift: 8,
// // BlueShift: 0,
// // })
// start := getNowMillisec()
// go func() {
// for {
// if getNowMillisec()-start >= 10000 {
// break
// }
// err = clientConn.FramebufferUpdateRequest(true, 0, 0, 1280, 800)
// if err != nil {
// logger.Errorf("error requesting fb update: %s", err)
// }
// time.Sleep(250 * time.Millisecond)
// }
// clientConn.Close()
// }()
for {
time.Sleep(time.Minute)

View File

@ -47,8 +47,8 @@ func main() {
TargetPort: *targetVncPort,
TargetPassword: *targetVncPass, //"vncPass",
ID: "dummySession",
//Status: SessionStatusActive,
//Type: SessionTypeRecordingProxy,
Status: proxy.SessionStatusInit,
Type: proxy.SessionTypeRecordingProxy,
}, // to be used when not using sessions
UsingSessions: false, //false = single session - defined in the var above
}

View File

@ -45,7 +45,7 @@ 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.SegmentMessageSeparator:
case common.SegmentMessageStart:
case common.SegmentRectSeparator:
case common.SegmentServerInitMessage:
serverInitMessage := seg.Message.(*common.ServerInit)

View File

@ -9,8 +9,9 @@ import (
"vncproxy/common"
"vncproxy/encodings"
"vncproxy/logger"
"vncproxy/player"
listeners "vncproxy/recorder"
"vncproxy/server"
listeners "vncproxy/tee-listeners"
)
type VncProxy struct {
@ -34,14 +35,11 @@ func (vp *VncProxy) createClientConnection(targetServerUrl string, vncPass strin
var noauth client.ClientAuthNone
authArr := []client.ClientAuth{&client.PasswordAuth{Password: vncPass}, &noauth}
//vncSrvMessagesChan := make(chan common.ServerMessage)
clientConn, err := client.NewClientConn(nc,
&client.ClientConfig{
Auth: authArr,
Exclusive: true,
})
//clientConn.Listener = split
if err != nil {
logger.Errorf("error creating client: %s", err)
@ -52,7 +50,7 @@ func (vp *VncProxy) createClientConnection(targetServerUrl string, vncPass strin
}
// if sessions not enabled, will always return the configured target server (only one)
func (vp *VncProxy) getTargetServerFromSession(sessionId string) (*VncSession, error) {
func (vp *VncProxy) getProxySession(sessionId string) (*VncSession, error) {
if !vp.UsingSessions {
if vp.SingleSession == nil {
@ -65,83 +63,95 @@ func (vp *VncProxy) getTargetServerFromSession(sessionId string) (*VncSession, e
func (vp *VncProxy) newServerConnHandler(cfg *server.ServerConfig, sconn *server.ServerConn) error {
recFile := "recording" + strconv.FormatInt(time.Now().Unix(), 10) + ".rbs"
recPath := path.Join(vp.RecordingDir, recFile)
rec, err := listeners.NewRecorder(recPath)
if err != nil {
logger.Errorf("Proxy.newServerConnHandler can't open recorder save path: %s", recPath)
return err
}
session, err := vp.getTargetServerFromSession(sconn.SessionId)
session, err := vp.getProxySession(sconn.SessionId)
if err != nil {
logger.Errorf("Proxy.newServerConnHandler can't get session: %d", sconn.SessionId)
return err
}
// for _, l := range rfbListeners {
// sconn.Listeners.AddListener(l)
// }
sconn.Listeners.AddListener(rec)
var rec *listeners.Recorder
if session.Type == SessionTypeRecordingProxy {
recFile := "recording" + strconv.FormatInt(time.Now().Unix(), 10) + ".rbs"
recPath := path.Join(vp.RecordingDir, recFile)
rec, err := listeners.NewRecorder(recPath)
if err != nil {
logger.Errorf("Proxy.newServerConnHandler can't open recorder save path: %s", recPath)
return err
}
//clientSplitter := &common.MultiListener{}
cconn, err := vp.createClientConnection(session.TargetHostname+":"+session.TargetPort, session.TargetPassword)
if err != nil {
logger.Errorf("Proxy.newServerConnHandler error creating connection: %s", err)
return err
}
cconn.Listeners.AddListener(rec)
//cconn.Listener = clientSplitter
//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)
//serverMsgRepeater := &listeners.WriteTo{sconn, "vnc-client-bound"}
//cconn.Listeners.AddListener(serverMsgRepeater)
// gets the messages from the server part (from vnc-client),
// and write through the client to the actual vnc-server
//clientMsgRepeater := &listeners.WriteTo{cconn, "vnc-server-bound"}
clientUpdater := &ClientUpdater{cconn}
sconn.Listeners.AddListener(clientUpdater)
err = cconn.Connect()
if err != nil {
logger.Errorf("Proxy.newServerConnHandler error connecting to client: %s", err)
return err
sconn.Listeners.AddListener(rec)
}
encs := []common.IEncoding{
&encodings.RawEncoding{},
&encodings.TightEncoding{},
&encodings.EncCursorPseudo{},
//encodings.TightPngEncoding{},
&encodings.RREEncoding{},
&encodings.ZLibEncoding{},
&encodings.ZRLEEncoding{},
&encodings.CopyRectEncoding{},
&encodings.CoRREEncoding{},
&encodings.HextileEncoding{},
session.Status = SessionStatusInit
if session.Type == SessionTypeProxyPass || session.Type == SessionTypeRecordingProxy {
cconn, err := vp.createClientConnection(session.TargetHostname+":"+session.TargetPort, 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.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
}
}
cconn.Encs = encs
//err = cconn.MsgSetEncodings(encs)
if err != nil {
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() {
//chServer := make(chan common.ClientMessage)
//chClient := make(chan common.ServerMessage)
secHandlers := []server.SecurityHandler{&server.ServerAuthNone{}}
if vp.ProxyVncPassword != "" {
@ -157,10 +167,6 @@ func (vp *VncProxy) StartListening() {
Width: uint16(1024),
NewConnHandler: vp.newServerConnHandler,
UseDummySession: !vp.UsingSessions,
// func(cfg *server.ServerConfig, conn *server.ServerConn) error {
// vp.newServerConnHandler(cfg, conn)
// return nil
// },
}
if vp.TcpListeningUrl != "" && vp.WsListeningUrl != "" {

View File

@ -16,7 +16,7 @@ func TestProxy(t *testing.T) {
TargetPort: "5903",
TargetPassword: "Ch_#!T@8",
ID: "dummySession",
Status: SessionStatusActive,
Status: SessionStatusInit,
Type: SessionTypeRecordingProxy,
}, // to be used when not using sessions
UsingSessions: false, //false = single session - defined in the var above

View File

@ -6,6 +6,7 @@ type SessionType int
const (
SessionStatusInit SessionStatus = iota
SessionStatusActive
SessionStatusError
)
const (
@ -21,4 +22,5 @@ type VncSession struct {
ID string
Status SessionStatus
Type SessionType
ReplayFilePath string
}

View File

@ -1,4 +1,4 @@
package listeners
package recorder
import (
"bytes"
@ -129,7 +129,7 @@ func (r *Recorder) HandleRfbSegment(data *common.RfbSegment) error {
}()
switch data.SegmentType {
case common.SegmentMessageSeparator:
case common.SegmentMessageStart:
if !r.sessionStartWritten {
logger.Debugf("Recorder.HandleRfbSegment: writing start session segment: %v", r.serverInitMessage)
r.writeStartSession(r.serverInitMessage)

47
recorder/rfb-requester.go Normal file
View File

@ -0,0 +1,47 @@
package recorder
import (
"time"
"vncproxy/client"
"vncproxy/common"
"vncproxy/logger"
)
type RfbRequester struct {
Conn *client.ClientConn
Name string
Width uint16
Height uint16
lastRequestTime time.Time
}
func (p *RfbRequester) Consume(seg *common.RfbSegment) error {
logger.Debugf("WriteTo.Consume ("+p.Name+"): got segment type=%s", seg.SegmentType)
switch seg.SegmentType {
case common.SegmentServerInitMessage:
serverInitMessage := seg.Message.(*common.ServerInit)
p.Conn.FrameBufferHeight = serverInitMessage.FBHeight
p.Conn.FrameBufferWidth = serverInitMessage.FBWidth
p.Conn.DesktopName = string(serverInitMessage.NameText)
p.Conn.SetPixelFormat(&serverInitMessage.PixelFormat)
p.Width = serverInitMessage.FBWidth
p.Height = serverInitMessage.FBHeight
p.lastRequestTime = time.Now()
p.Conn.FramebufferUpdateRequest(false, 0, 0, p.Width, p.Height)
case common.SegmentMessageStart:
case common.SegmentRectSeparator:
case common.SegmentBytes:
case common.SegmentFullyParsedClientMessage:
case common.SegmentMessageEnd:
// minTimeBetweenReq := 300 * time.Millisecond
// timeForNextReq := p.lastRequestTime.Unix() + minTimeBetweenReq.Nanoseconds()/1000
// if seg.UpcomingObjectType == int(common.FramebufferUpdate) && time.Now().Unix() > timeForNextReq {
//time.Sleep(300 * time.Millisecond)
p.Conn.FramebufferUpdateRequest(true, 0, 0, p.Width, p.Height)
//}
default:
}
return nil
}

View File

@ -1,9 +1,8 @@
# TODO:
* add replay flow to proxy
* set correct status for each flow in proxy
* create 2 cmdline apps (recorder & proxy) - proxy will also replay (depending on session type & cmdline flags)
* add replay logics to proxy (depending on session type & cmdline flags)
* set correct status for each flow in proxy (replay / prox+record / prox / ..)
* improve recorder so it will save RFB response before sending another RFB update request
* code stuff:
* move encodings to be on the framebufferupdate message object
* clear all messages read functions from updating stuff, move modification logic to another listener