github.com/amitbet/vnc2video@v0.0.0-20190616012314-9d50b9dab1d9/example/proxy/main.go (about) 1 package main 2 3 import ( 4 "bytes" 5 "context" 6 "encoding/base64" 7 "fmt" 8 "io" 9 "net" 10 "net/http" 11 _ "net/http/pprof" 12 "net/url" 13 "strings" 14 "sync" 15 "time" 16 vnc "github.com/amitbet/vnc2video" 17 "github.com/amitbet/vnc2video/logger" 18 ) 19 20 type Auth struct { 21 Username []byte 22 Password []byte 23 } 24 25 type Proxy struct { 26 cc vnc.Conn 27 conns chan vnc.Conn 28 inp chan vnc.ClientMessage 29 out chan vnc.ServerMessage 30 } 31 32 var ( 33 cliconns = make(map[string]*Proxy) 34 srvconns = make(map[vnc.Conn]string) 35 m sync.Mutex 36 ) 37 38 func newConn(hostport string, password []byte) (vnc.Conn, chan vnc.ClientMessage, chan vnc.ServerMessage, chan vnc.Conn, error) { 39 fmt.Printf("new conn to %s with %s\n", hostport, password) 40 if cc, ok := cliconns[hostport]; ok { 41 return cc.cc, cc.inp, cc.out, cc.conns, nil 42 } 43 c, err := net.DialTimeout("tcp", hostport, 10*time.Second) 44 if err != nil { 45 return nil, nil, nil, nil, err 46 } 47 cchServer := make(chan vnc.ServerMessage) 48 cchClient := make(chan vnc.ClientMessage) 49 errorCh := make(chan error) 50 ccfg := &vnc.ClientConfig{ 51 SecurityHandlers: []vnc.SecurityHandler{&vnc.ClientAuthVNC{Password: password}}, 52 PixelFormat: vnc.PixelFormat32bit, 53 ClientMessageCh: cchClient, 54 ServerMessageCh: cchServer, 55 //ServerMessages: vnc.DefaultServerMessages, 56 Encodings: []vnc.Encoding{&vnc.RawEncoding{}}, 57 ErrorCh: errorCh, 58 } 59 csrv := make(chan vnc.Conn) 60 inp := make(chan vnc.ClientMessage) 61 out := make(chan vnc.ServerMessage) 62 fmt.Printf("connect to vnc\n") 63 cc, err := vnc.Connect(context.Background(), c, ccfg) 64 if err != nil { 65 return nil, nil, nil, nil, err 66 } 67 fmt.Printf("connected to vnc %#+v\n", cc) 68 ds := &vnc.DefaultClientMessageHandler{} 69 go ds.Handle(cc) 70 go handleIO(cc, inp, out, csrv) 71 72 return cc, inp, out, csrv, nil 73 } 74 75 func handleIO(cli vnc.Conn, inp chan vnc.ClientMessage, out chan vnc.ServerMessage, csrv chan vnc.Conn) { 76 fmt.Printf("handle io\n") 77 ccfg := cli.Config().(*vnc.ClientConfig) 78 defer cli.Close() 79 var conns []vnc.Conn 80 //var prepared bool 81 82 for { 83 select { 84 case err := <-ccfg.ErrorCh: 85 for _, srv := range conns { 86 srv.Close() 87 } 88 fmt.Printf("err %v\n", err) 89 return 90 case msg := <-ccfg.ServerMessageCh: 91 for _, srv := range conns { 92 scfg := srv.Config().(*vnc.ServerConfig) 93 scfg.ServerMessageCh <- msg 94 } 95 case msg := <-inp: 96 // messages from real clients 97 fmt.Printf("3 %#+v\n", msg) 98 switch msg.Type() { 99 case vnc.SetPixelFormatMsgType: 100 101 case vnc.SetEncodingsMsgType: 102 var encTypes []vnc.EncodingType 103 encs := []vnc.Encoding{ 104 // &vnc.TightPngEncoding{}, 105 &vnc.CopyRectEncoding{}, 106 &vnc.RawEncoding{}, 107 } 108 for _, senc := range encs { 109 for _, cenc := range msg.(*vnc.SetEncodings).Encodings { 110 if cenc == senc.Type() { 111 encTypes = append(encTypes, senc.Type()) 112 } 113 } 114 } 115 ccfg.ClientMessageCh <- &vnc.SetEncodings{Encodings: encTypes} 116 default: 117 ccfg.ClientMessageCh <- msg 118 } 119 case msg := <-out: 120 fmt.Printf("4 %#+v\n", msg) 121 case srv := <-csrv: 122 conns = append(conns, srv) 123 } 124 125 } 126 127 } 128 129 type HijackHandler struct{} 130 131 func (*HijackHandler) Handle(c vnc.Conn) error { 132 m.Lock() 133 defer m.Unlock() 134 hostport, ok := srvconns[c] 135 if !ok { 136 return fmt.Errorf("client connect in server pool not found") 137 } 138 proxy, ok := cliconns[hostport] 139 if !ok { 140 return fmt.Errorf("client connect to qemu not found") 141 } 142 cfg := c.Config().(*vnc.ServerConfig) 143 cfg.ClientMessageCh = proxy.inp 144 cfg.ServerMessageCh = proxy.out 145 146 proxy.conns <- c 147 ds := &vnc.DefaultServerMessageHandler{} 148 go ds.Handle(c) 149 return nil 150 } 151 152 type AuthVNCHTTP struct { 153 c *http.Client 154 vnc.ServerAuthVNC 155 } 156 157 func (auth *AuthVNCHTTP) Auth(c vnc.Conn) error { 158 auth.ServerAuthVNC.Challenge = []byte("clodo.ruclodo.ru") 159 if err := auth.ServerAuthVNC.WriteChallenge(c); err != nil { 160 return err 161 } 162 if err := auth.ServerAuthVNC.ReadChallenge(c); err != nil { 163 return err 164 } 165 166 buf := new(bytes.Buffer) 167 enc := base64.NewEncoder(base64.StdEncoding, buf) 168 enc.Write(auth.ServerAuthVNC.Crypted) 169 enc.Close() 170 171 v := url.Values{} 172 v.Set("hash", buf.String()) 173 buf.Reset() 174 src, _, _ := net.SplitHostPort(c.Conn().RemoteAddr().String()) 175 v.Set("ip", src) 176 res, err := auth.c.PostForm("https://api.ix.clodo.ru/system/vnc", v) 177 if err != nil { 178 return err 179 } 180 if res.StatusCode != 200 || res.Body == nil { 181 if res.Body != nil { 182 io.Copy(buf, res.Body) 183 } 184 fmt.Printf("failed to get auth data: code %d body %s\n", res.StatusCode, buf.String()) 185 defer buf.Reset() 186 return fmt.Errorf("failed to get auth data: code %d body %s", res.StatusCode, buf.String()) 187 } 188 _, err = io.Copy(buf, res.Body) 189 if err != nil { 190 return fmt.Errorf("failed to get auth data: %s", err.Error()) 191 } 192 logger.Debugf("http auth: %s\n", buf.Bytes()) 193 res.Body.Close() 194 data := strings.Split(buf.String(), " ") 195 if len(data) < 2 { 196 return fmt.Errorf("failed to get auth data data invalid") 197 } 198 buf.Reset() 199 200 hostport := string(data[0]) 201 password := []byte(data[1]) 202 203 m.Lock() 204 defer m.Unlock() 205 cc, inp, out, conns, err := newConn(hostport, password) 206 if err != nil { 207 return err 208 } 209 cliconns[hostport] = &Proxy{cc, conns, inp, out} 210 srvconns[c] = hostport 211 c.SetWidth(cc.Width()) 212 c.SetHeight(cc.Height()) 213 return nil 214 } 215 216 func (*AuthVNCHTTP) Type() vnc.SecurityType { 217 return vnc.SecTypeVNC 218 } 219 220 func (*AuthVNCHTTP) SubType() vnc.SecuritySubType { 221 return vnc.SecSubTypeUnknown 222 } 223 224 func main() { 225 go func() { 226 logger.Info(http.ListenAndServe(":6060", nil)) 227 }() 228 229 ln, err := net.Listen("tcp", ":6900") 230 if err != nil { 231 logger.Fatalf("Error listen. %v", err) 232 } 233 234 schClient := make(chan vnc.ClientMessage) 235 schServer := make(chan vnc.ServerMessage) 236 237 scfg := &vnc.ServerConfig{ 238 SecurityHandlers: []vnc.SecurityHandler{ 239 &AuthVNCHTTP{c: &http.Client{}}, 240 }, 241 Encodings: []vnc.Encoding{ 242 // &vnc.TightPngEncoding{}, 243 &vnc.CopyRectEncoding{}, 244 &vnc.RawEncoding{}, 245 }, 246 PixelFormat: vnc.PixelFormat32bit, 247 ClientMessageCh: schClient, 248 ServerMessageCh: schServer, 249 //ClientMessages: vnc.DefaultClientMessages, 250 DesktopName: []byte("vnc proxy"), 251 } 252 scfg.Handlers = append(scfg.Handlers, vnc.DefaultServerHandlers...) 253 scfg.Handlers = append(scfg.Handlers[:len(scfg.Handlers)-1], &HijackHandler{}) 254 vnc.Serve(context.Background(), ln, scfg) 255 }