Golang websocket.Conn类代码示例

func remoteControl(ws *websocket.Conn) {
	pptRequest := &pptctrl.Request{Response: make(chan string)}

	defer ws.Close()
	defer close(pptRequest.Response)

	r := bufio.NewReader(ws)

	for {
		data, err := r.ReadBytes('\n')
		fmt.Printf("%s", data)

		if err != nil {
			fmt.Printf("Error occured: %s\n", err.Error())
			break
		}

		switch data[0] {
		case '!': // PowerPoint control
			pptRequest.Data = data

			pptctrl.SendRequest(pptRequest)

			// block current goroutine
			extraReturnInfo := <-pptRequest.Response

			if len(extraReturnInfo) > 0 {
				sendPowerPointExtraReturnInfo(data, extraReturnInfo, ws)
			}
		}
	}
}

//websocket main handler
func wsMain(ws *websocket.Conn) {
	defer ws.Close()

	wsCtx := newWSContext(ws)

	var msg string
	for {
		//read message
		err := websocket.Message.Receive(ws, &msg)
		if err != nil {
			break
		}
		log.Debug("[wsMain]endpoint's msg = " + msg)

		//parse message
		wsMsg, err := NewEndptMsgFromStr(msg)

		if err != nil {
			log.Error("[wsMain]failed to unmarshal json :" + err.Error())
			continue
		}
		wsCtx.UserID = wsMsg.UserID

		if wsMsg.Domain == "irc" && wsCtx.LoggedIn {
			ClientDoIRCCmd(wsMsg, ws)
		} else {
			dispatchBoksHandler(wsCtx, wsMsg)
		}
	}

	if wsCtx.LoggedIn {
		UserLogout(wsCtx.UserID, ws)
	}
	log.Debug("[wsMain]endpoint exited")
}

func readMsg(conn *websocket.Conn, owner string) {
	tmp := make([]byte, 32)
	buf := make([]byte, 0)
	for {
		readlen, err := conn.Read(tmp)
		if err != nil {
			break
		}

		buf = append(buf, tmp[0:readlen]...)
		lensep := strings.Index(string(buf), ":")
		if lensep < 0 {
			continue
		}

		msglen, err := strconv.Atoi(string(buf[0:lensep]))
		if err != nil {
			log.Println("error: ", err)
			break
		}

		if len(buf) < msglen+lensep+1 {
			continue
		}

		msg := Msg{owner, make([]byte, msglen)}
		copy(msg.data, buf[lensep+1:msglen+lensep+1])
		read <- msg
		buf = buf[lensep+msglen+1:]
	}
}

func wsHandler(ws *wbs.Conn) {
	defer ws.Close()

	log.Println("Websocket connection recieved.")
	log.Println("Handeling user.")
	handleUser(ws)
}

func (fh *FakeHandler) handle(conn *websocket.Conn) {
	fh.call()
	request := conn.Request()
	fh.setLastURL(request.URL.String())
	fh.setAuthHeader(request.Header.Get("Authorization"))

	if fh.messageReceived != nil {
		go func() {
			for {
				buffer := make([]byte, 1024)
				_, err := conn.Read(buffer)

				if err == nil {
					fh.messageReceived <- true
				} else {
					break
				}
			}
		}()
	}

	for _, protoMessage := range fh.Messages {
		if protoMessage == nil {
			conn.Write([]byte{})
		} else {
			message, err := proto.Marshal(protoMessage)
			Expect(err).ToNot(HaveOccurred())

			conn.Write(message)
		}
	}

	<-fh.closeConnection
	conn.Close()
}

func connect(conn *websocket.Conn) {
	logger.Printf("websocket from %s", conn.RemoteAddr())
	defer conn.Close()
	var form []byte
	var f Form
	if err := websocket.Message.Receive(conn, &form); err != nil {
		return
	}
	if err := json.Unmarshal(form, &f); err != nil {
		loggedmessage(conn, "invalid request: %s (%s)", form, err)
		return
	}
	loggedmessage(conn, "opening connection to %s for %s", f.Serv, f.Nick)
	client, err := NewClient(f.Serv, f.Nick)
	if err != nil {
		websocket.Message.Send(conn, "connection error: "+err.Error())
		return
	}
	defer func() {
		logger.Printf("closing connection to %s for %s", f.Serv, f.Nick)
		websocket.Message.Send(conn, "connection closed.")
		client.Cmd("QUIT :%s", "client left.")
		client.Close()
	}()
	logger.Printf("joining channel %s", f.Chan)
	client.Cmd("JOIN %s", f.Chan)
	for line := range client.lines {
		// send {"system": message} or {"line": message}
		websocket.JSON.Send(conn, line)
	}
}

func (proxy *Proxy) HandleWebSocket(clientWS *websocket.Conn) {
	req := clientWS.Request()
	req.ParseForm()
	req.Form.Get("app")
	clientAddress := clientWS.RemoteAddr()

	appId := req.Form.Get("app")

	extractAuthTokenFromUrl := func(u *url.URL) string {
		authorization := ""
		queryValues := u.Query()
		if len(queryValues["authorization"]) == 1 {
			authorization = queryValues["authorization"][0]
		}
		return authorization
	}

	authToken := clientWS.Request().Header.Get("Authorization")
	if authToken == "" {
		authToken = extractAuthTokenFromUrl(req.URL)
	}

	if authorized, errorMessage := proxy.isAuthorized(appId, authToken, clientAddress); !authorized {
		data, err := proto.Marshal(errorMessage)
		if err != nil {
			proxy.logger.Errorf("Error marshalling log message: %s", err)
		}
		websocket.Message.Send(clientWS, data)
		clientWS.Close()
		return
	}

	defer clientWS.Close()

	proxy.logger.Debugf("Output Proxy: Request for app: %v", req.Form.Get("app"))
	serverWSs := make([]*websocket.Conn, len(proxy.hashers))
	for index, hasher := range proxy.hashers {
		proxy.logger.Debugf("Output Proxy: Servers in group [%v]: %v", index, hasher.LoggregatorServers())

		server := hasher.GetLoggregatorServerForAppId(appId)
		proxy.logger.Debugf("Output Proxy: AppId is %v. Using server: %v", appId, server)

		config, err := websocket.NewConfig("ws://"+server+req.URL.RequestURI(), "http://localhost")

		if err != nil {
			proxy.logger.Errorf("Output Proxy: Error creating config for websocket - %v", err)
		}

		serverWS, err := websocket.DialConfig(config)
		if err != nil {
			proxy.logger.Errorf("Output Proxy: Error connecting to loggregator server - %v", err)
		}

		if serverWS != nil {
			serverWSs[index] = serverWS
		}
	}
	proxy.forwardIO(serverWSs, clientWS)

}

func newClient(ws *websocket.Conn, h *hub, srv *wsServer) *wsClient {
	if ws == nil {
		log.Panicln("ws cannot be nil")
	}

	if srv == nil {
		log.Panicln("server cannot be nil")
	}

	maxID++

	return &wsClient{
		id:   maxID,
		ws:   ws,
		h:    h,
		srv:  srv,
		subs: make(map[string]string),

		inMsgChan:  make(chan *wsMessage),
		outMsgChan: make(chan *wsMessage),
		doneChan:   make(chan bool),

		socketID: sha1Sum(fmt.Sprintf("%s-%s-%d", ws.RemoteAddr().String(),
			ws.LocalAddr().String(), maxID)),
	}
}

// Starting point of websocket connection
//	*	Verify identity
//	*	Register send/recieve channel
//	* 	Manage send/recieve for duration of connection
func (t *Server) HandleWebsocket(conn *websocket.Conn) {
	defer conn.Close() //Close connection at end of this function

	//Register Connection
	c, err := t.registerConnection(conn)
	if err != nil {
		log.Error("postmaster: error registering connection: %s", err)
		return
	}

	//Setup goroutine to send all message on chan
	go func() {
		for msg := range c.out {
			log.Trace("postmaster: sending message: %s", msg)
			err := websocket.Message.Send(conn, msg)
			if err != nil {
				log.Error("postmaster: error sending message: %s", err)
			}
		}
	}()

	//Setup message recieving (Blocking for life of connection)
	t.recieveOnConn(c, conn)

	//Call disconnection method
	//FIXME Figure out why pendingAuth is nil sometimes
	if t.OnDisconnect != nil && c.pendingAuth != nil {
		t.OnDisconnect(c.pendingAuth.authKey, c.pendingAuth.authExtra)
	}

	//Unregister connection
	delete(t.connections, c.id)
}

// Connection handler. This function called after new client
// connected to websocket server.
// Also this method performs register user - client must send valid apiKey
// to receive messages from redis.
func wsHandler(ws *websocket.Conn) {
	log.Println("New client")
	defer ws.Close()
	// websocket.Message.Send(ws, "Hello dear user!")

	for {
		// Message received from client
		var message string

		// Read messages from client
		// Code blocks here, after any message received
		// will resume execution.
		if err := websocket.Message.Receive(ws, &message); err != nil {
			log.Println("Error receiving message. Closing connection.")
			return
		}

		// Register user
		// TODO: Cache user
		user, _ := common.FindUserByApiKey(message)
		if user != nil {
			log.Println("Registering apiKey", user.ApiKey)
			clients[user.ApiKey] = ws
		} else {
			log.Println("Error registering user", message)
		}
	}
}

func (gs *gameServer) clientListenRead(ws *websocket.Conn) {
	defer func() {
		ws.Close()
	}()

	for {
		select {
		default:
			var move util.ClientMove
			err := websocket.JSON.Receive(ws, &move)
			if err == io.EOF {
				return
				// EOF!
			} else if err != nil {
				LOGE.Println(err)
			} else {
				var dir lib2048.Direction
				switch move.Direction {
				case 0:
					dir = lib2048.Up
				case 1:
					dir = lib2048.Right
				case 2:
					dir = lib2048.Down
				case 3:
					dir = lib2048.Left
				}
				LOGV.Println("Received", dir, "from web client")
				move := lib2048.NewMove(dir)
				gs.clientMoveCh <- move
			}
		}
	}
}

// WebSocket server to handle chat between clients
func SockServer(ws *websocket.Conn) {
	var err error
	var clientMessage string
	// use []byte if websocket binary type is blob or arraybuffer
	// var clientMessage []byte

	// cleanup on server side
	defer func() {
		if err = ws.Close(); err != nil {
			log.Println("Websocket could not be closed", err.Error())
		}
	}()

	client := ws.Request().RemoteAddr
	log.Println("Client connected:", client)
	sockCli := ClientConn{ws, client}
	ActiveClients[sockCli] = 0
	log.Println("Number of clients connected ...", len(ActiveClients))

	updateNewClient(sockCli)
	// for loop so the websocket stays open otherwise
	// it'll close after one Receive
	for {
		if err = Message.Receive(ws, &clientMessage); err != nil {
			// If we cannot Read then the connection is closed
			log.Println("Websocket Disconnected waiting", err.Error())
			// remove the ws client conn from our active clients
			delete(ActiveClients, sockCli)
			log.Println("Number of clients still connected ...", len(ActiveClients))
			return
		}
		piano_ctrl <- clientMessage
	}
}

func transmitter(id string, ws *websocket.Conn, c chan *registrar.T, ctx context.Context, cancel context.CancelFunc) {
	var err error
	var data *registrar.T
	defer ws.Close()
	//defer close(c)
	defer cancel()
	defer registrar.RemoveConnection(id)
Loop:
	for {
		select {
		case data = <-c:
			err = websocket.JSON.Send(ws, *data)
			//websocket.Message.Send(ws, data.Msg)
			if err != nil {
				if !ws.IsClientConn() {
					log.Printf("transmitter closed\n")
				} else {
					log.Printf("transmitter error %v\n", err)
				}
				break Loop
			}
		case <-ctx.Done():
			log.Printf("transmitter closed")
			break Loop
		}
	}
}

//Function that handles all the request and creates response
func WShandler(ws *websocket.Conn) {
	var err error

	// cleanout
	defer ws.Close()
	//Here we are creating list of clients that gets connected
	socketClientIP := ws.Request().RemoteAddr
	socketClient := Clients{ws, socketClientIP}
	ActiveClients[socketClient] = 0
	log.Println("Total clients live:", len(ActiveClients))
	//For loop to keep it open, It closes after one Send/Recieve
	for {
		var reply string
		var clientMessage string
		if err = websocket.Message.Receive(ws, &reply); err != nil {
			log.Println("Can't receive ", socketClientIP, err.Error())
		}
		fmt.Println("Received back from client: " + reply)

		clientMessage = socketClient.IP + " : " + reply
		//ForEACH client conected, send back the msg to everyone
		for cs, _ := range ActiveClients {
			if err = websocket.Message.Send(cs.websocket, clientMessage); err != nil {
				// It could not send message to a peer
				log.Println("Could not send message to ", cs.IP, err.Error())
			}
		}
	}
}

func handleConnection(conn *net.TCPConn) error {
	defer conn.Close()

	handlerChan <- 1
	defer func() {
		handlerChan <- -1
	}()

	var ws *websocket.Conn

	conn.SetDeadline(time.Now().Add(socksTimeout))
	err := AwaitSocks4aConnect(conn, func(dest string) (*net.TCPAddr, error) {
		// Disable deadline.
		conn.SetDeadline(time.Time{})
		Log("SOCKS request for %s", dest)
		destAddr, err := net.ResolveTCPAddr("tcp", dest)
		if err != nil {
			return nil, err
		}
		wsUrl := url.URL{Scheme: "ws", Host: dest}
		ws, err = websocket.Dial(wsUrl.String(), "", wsUrl.String())
		if err != nil {
			return nil, err
		}
		Log("WebSocket connection to %s", ws.Config().Location.String())
		return destAddr, nil
	})
	if err != nil {
		return err
	}
	defer ws.Close()
	proxy(conn, ws)
	return nil
}