// handleWebsocket connection. Update to func httpInvoke(w http.ResponseWriter, r *http.Request) { ctx := remote.NewConnContext() writeEvents := func(evs []event.Event, buf *bytes.Buffer) error { if len(evs) > 0 { buf.Reset() for _, ev := range evs { if nil != ev { event.EncryptEvent(buf, ev, &ctx.CryptoContext) } } if buf.Len() > 0 { _, err := w.Write(buf.Bytes()) if nil == err { w.(http.Flusher).Flush() } return err } } return nil } //reqbuf := readRequestBuffer(r) var wbuf bytes.Buffer ress, err := handleRequestBody(r, ctx) if nil != err { log.Printf("[ERROR]connection %s:%d error:%v with path:%s ", ctx.User, ctx.ConnIndex, err, r.URL.Path) w.WriteHeader(400) } else { w.WriteHeader(200) if strings.HasSuffix(r.URL.Path, "pull") { begin := time.Now() period, _ := strconv.Atoi(r.Header.Get("X-PullPeriod")) if period <= 0 { period = 15 } writeEvents(ress, &wbuf) queue := remote.GetEventQueue(ctx.ConnId, true) defer remote.ReleaseEventQueue(queue) for { if time.Now().After(begin.Add(time.Duration(period) * time.Second)) { log.Printf("Stop puller after %ds for conn:%d", period, ctx.ConnIndex) break } evs, err := queue.PeekMulti(2, 1*time.Millisecond, false) if nil != err { continue } err = writeEvents(evs, &wbuf) if nil != err { log.Printf("HTTP write error:%v", err) return } queue.DiscardPeeks(false) } } } }
func serveProxyConn(conn net.Conn, vs *vpsServer) { if nil != vs { atomic.AddInt32(&vs.aliveConns, 1) } atomic.AddInt32(&totalConn, 1) bufconn := bufio.NewReader(conn) deferFunc := func() { conn.Close() atomic.AddInt32(&totalConn, -1) if nil != vs { atomic.AddInt32(&vs.aliveConns, -1) } } defer deferFunc() ctx := remote.NewConnContext() writeEvents := func(evs []event.Event, buf *bytes.Buffer) error { if len(evs) > 0 { buf.Reset() for _, ev := range evs { if nil != ev { event.EncryptEvent(buf, ev, &ctx.CryptoContext) } } if buf.Len() > 0 { conn.SetWriteDeadline(time.Now().Add(15 * time.Second)) b := buf.Bytes() _, err := conn.Write(b) return err } } return nil } var rbuf bytes.Buffer var wbuf bytes.Buffer writeTaskRunning := false connClosed := false reader := &helper.BufferChunkReader{bufconn, nil} for !connClosed { conn.SetReadDeadline(time.Now().Add(60 * time.Second)) rbuf.Grow(8192) rbuf.ReadFrom(reader) if nil != reader.Err { conn.Close() connClosed = true } ress, err := remote.HandleRequestBuffer(&rbuf, ctx) if nil != err { if err != event.EBNR { log.Printf("[ERROR]connection %s:%d error:%v", ctx.User, ctx.ConnIndex, err) conn.Close() connClosed = true return } } else { writeEvents(ress, &wbuf) if !writeTaskRunning && len(ctx.User) > 0 && ctx.ConnIndex >= 0 { writeTaskRunning = true go func() { var lastEventTime time.Time queue := remote.GetEventQueue(ctx.ConnId, true) for !connClosed { evs, err := queue.PeekMulti(10, 1*time.Millisecond, false) now := time.Now() if ctx.Closing { evs = []event.Event{&event.ChannelCloseACKEvent{}} if remote.ServerConf.MaxDynamicPort > 0 { nvs := selectDynamicVPServer() if nil != nvs { evs = append(evs, &event.PortUnicastEvent{Port: nvs.port}) } } } else { if nil != err { if remote.GetSessionTableSize() > 0 && lastEventTime.Add(10*time.Second).Before(now) { evs = []event.Event{event.NewHeartBeatEvent()} } else { continue } } } err = writeEvents(evs, &wbuf) if nil != err { log.Printf("TCP write error####:%v %d", err, len(evs)) } else { queue.DiscardPeeks(false) } if ctx.Closing { break } lastEventTime = now if nil != err { log.Printf("TCP write error:%v", err) conn.Close() break } else { //queue.DiscardPeeks(false) } } remote.ReleaseEventQueue(queue) }() } } } }
// handleWebsocket connection. Update to func websocketInvoke(w http.ResponseWriter, r *http.Request) { if r.Method != "GET" { http.Error(w, "Method not allowed", 405) return } ws, err := upgrader.Upgrade(w, r, nil) if err != nil { //log.WithField("err", err).Println("Upgrading to websockets") http.Error(w, "Error Upgrading to websockets", 400) return } ctx := remote.NewConnContext() writeEvents := func(evs []event.Event, wbuf *bytes.Buffer) error { if len(evs) > 0 { //var buf bytes.Buffer wbuf.Reset() for _, ev := range evs { if nil != ev { event.EncryptEvent(wbuf, ev, &ctx.CryptoContext) } } if wbuf.Len() > 0 { return ws.WriteMessage(websocket.BinaryMessage, wbuf.Bytes()) } return nil } return nil } //log.Printf("###Recv websocket connection") buf := bytes.NewBuffer(nil) var wbuf bytes.Buffer wsClosed := false var queue *remote.ConnEventQueue for { mt, data, err := ws.ReadMessage() if err != nil { if err != io.EOF { log.Printf("Websoket read error:%v", err) } wsClosed = true break } switch mt { case websocket.BinaryMessage: if buf.Len() == 0 { buf = bytes.NewBuffer(data) } else { buf.Write(data) } ress, err := remote.HandleRequestBuffer(buf, ctx) if nil != err { log.Printf("[ERROR]connection %s:%d error:%v", ctx.User, ctx.ConnIndex, err) ws.Close() wsClosed = true } else { writeEvents(ress, &wbuf) if nil == queue && len(ctx.User) > 0 && ctx.ConnIndex >= 0 { queue = remote.GetEventQueue(ctx.ConnId, true) go func() { var wwbuf bytes.Buffer for !wsClosed { evs, err := queue.PeekMulti(2, 1*time.Millisecond, false) if ctx.Closing { evs = []event.Event{&event.ChannelCloseACKEvent{}} } else { if nil != err { continue } } err = writeEvents(evs, &wwbuf) if ctx.Closing { break } if nil != err { log.Printf("Websoket write error:%v", err) break } else { queue.DiscardPeeks(false) } } remote.ReleaseEventQueue(queue) }() } } default: log.Printf("Invalid websocket message type") ws.Close() } } wsClosed = true log.Printf("Close websocket connection:%d", ctx.ConnIndex) //ws.WriteMessage(websocket.CloseMessage, []byte{}) }