Golang Conn.Close方法代码示例

func (handler *handler) proxyConnectionTo(server *websocket.Conn) {
	handler.logger.Debugf("Output Proxy: Starting to listen to server %v", server.RemoteAddr().String())

	var logMessage []byte
	defer server.Close()
	count := 0
	for {
		_, data, err := server.ReadMessage()

		if err != nil {
			handler.logger.Errorf("Output Proxy: Error reading from the server - %v - %v", err, server.RemoteAddr().String())
			handler.Done()
			return
		}

		handler.logger.Debugf("Output Proxy: Got message from server %v bytes", len(logMessage))

		count++

		err = handler.writeMessage(data)

		if err != nil {
			handler.logger.Errorf("Output Proxy: Error writing to client websocket - %v", err)
			return
		}
	}
}

func (ws *WebSocker) writer(conn *websocket.Conn, user string, tmpl *template.Template) {
	pingTicker := time.NewTicker(pingPeriod)
	chl, chq := notifier.listen(user)
	ws.log.Println(user, "connected")
	defer func() {
		pingTicker.Stop()
		conn.Close()
		close(chq)
		ws.log.Println(user, "disconnected")
	}()
	for {
		var buf bytes.Buffer
		select {
		case p := <-chl:
			conn.SetWriteDeadline(time.Now().Add(writeWait))
			buf.Reset()
			if err := tmpl.Execute(&buf, p); err != nil {
				ws.log.Println(err)
				return
			}
			if err := conn.WriteMessage(websocket.TextMessage, buf.Bytes()); err != nil {
				return
			}
			ws.log.Println("notified:", user)
		case <-pingTicker.C:
			conn.SetWriteDeadline(time.Now().Add(writeWait))
			if err := conn.WriteMessage(websocket.PingMessage, []byte{}); err != nil {
				ws.log.Println("ping:", err)
				return
			}
		}
	}
}

func (bot *Bot) loop(conn *websocket.Conn) bool {
	defer conn.Close()
	for {
		messageType, bytes, err := conn.ReadMessage()
		if err != nil {
			// ReadMessage returns an error if the connection is closed
			return false
		}
		if messageType == websocket.BinaryMessage {
			continue // ignore binary messages
		}
		event, err := unpackJSON(bytes)
		if err != nil {
			log.WithFields(log.Fields{
				"raw bytes": bytes,
				"error":     err,
			}).Warn("message could not be unpacked")
			continue
		}
		log.WithFields(log.Fields{
			"event": event,
		}).Info("received event")
		eventType, ok := event["type"]
		if ok && eventType.(string) == "team_migration_started" {
			return true
		}
		wrappers := bot.handle(event)
		closeConnection := sendResponses(wrappers, conn)
		if closeConnection {
			return false
		}
	}
}

func ptySetupWs(ws *websocket.Conn, sizeFlag string) {
	ptym, cmd := start()
	setPtySize(ptym, sizeFlag)

	go func() {
		handleOutputWs(ptym, ws)
	}()

	go func() {
		handleInputWs(ptym, ws)
	}()

	// Listen for a new winsize on stdin.
	for {
		var newSize string
		_, scanErr := fmt.Scanln(&newSize)
		if scanErr != nil {
			fmt.Println("scan failed: ", scanErr)
		}

		setPtySize(ptym, newSize)
		fmt.Println("new size: ", newSize)
	}

	stop(ptym, cmd)
	ws.Close()
}

func (client *Client) ping(ws *websocket.Conn, addr string) {
	log.WithField("port", client.Port).Infoln("DailTLS ok: " + addr)
	info, err := client.getServerInfo()
	if err != nil {
		log.WithField("port", client.Port).Errorln("getServerInfo", err)
		return
	}

	ticker := time.NewTicker(info.PingSecond * time.Second)
	defer func() {
		ticker.Stop()
		ws.Close()
		log.WithField("port", client.Port).Infoln("Ws closed")
	}()
	log.WithField("port", client.Port).Infoln("Ws started")

	req := client.innerRequest("HEAD", HOST_OK)

	for {
		select {
		case <-ticker.C:
			ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second)
			res, err := client.h2Transport.RoundTrip(req.WithContext(ctx))
			if err != nil || res.StatusCode != http.StatusOK {
				cancel()
				return
			}
			cancel()
		}
	}
}

func PushMsg(ws *websocket.Conn, WSWriter chan []uint8) {

	defer ws.Close()

	for {
		msg := <-WSWriter

		if msg == nil {
			continue
		}

		// write message to client
		if err := ws.WriteMessage(websocket.TextMessage, []byte(msg)); err != nil {
			log.Println("[ErrorInfo]", err.Error())
			break
		}

		if string(msg) == "bye" {
			break
		}
	}

	endChan <- true

}

func (r *RancherService) pipeLogs(container *rancherClient.Container, conn *websocket.Conn) {
	defer conn.Close()

	log_name := strings.TrimPrefix(container.Name, r.context.ProjectName+"_")
	logger := r.context.LoggerFactory.Create(log_name)

	for {
		messageType, bytes, err := conn.ReadMessage()
		if messageType != websocket.TextMessage {
			continue
		}

		if err == io.EOF {
			return
		} else if err != nil {
			logrus.Errorf("Failed to read log: %v", err)
			return
		}

		if len(bytes) <= 3 {
			continue
		}

		if bytes[len(bytes)-1] != '\n' {
			bytes = append(bytes, '\n')
		}
		message := bytes[3:]

		if "01" == string(bytes[:2]) {
			logger.Out(message)
		} else {
			logger.Err(message)
		}
	}
}

func (np *NatsProxy) activateWSProxySubject(conn *websocket.Conn, wsID string) {
	np.addToWSMapper(conn, wsID)
	np.conn.Subscribe("WS_OUT"+wsID, func(m *nats.Msg) {
		err := conn.WriteMessage(websocket.TextMessage, m.Data)
		if err != nil {
			log.Println("Error writing a message", err)
		}
	})
	go func() {
		for {
			if _, p, err := conn.ReadMessage(); err == nil {
				np.conn.Publish("WS_IN"+wsID, p)
			} else {
				np.removeFromWSMapper(conn, wsID)
				conn.Close()
				// If websocket is closed normally RFC6455
				// code 1000, then no error logged
				if !websocket.IsCloseError(err, websocket.CloseNormalClosure) {
					logWebsocketError(wsID, err)
				}
				break
			}
		}
	}()
}

// writer runs in a goroutine for each connected WS client. It emits all message returned by the observer.
func writer(ws *websocket.Conn, a *ApiHandlers) {
	pingTicker := time.NewTicker(pingPeriod)
	s := a.getDBSubscriber()
	jww.INFO.Println("Opened WebSocket connection.")
	defer func(is *subscriber) {
		jww.INFO.Println("Closing WebSocket connection.")
		is.quitChan <- true
		pingTicker.Stop()
		ws.Close()
	}(s)

	for {
		select {
		case msg := <-s.bufChan:
			ws.SetWriteDeadline(time.Now().Add(writeWait))
			if err := ws.WriteJSON(msg); err != nil {
				return
			}
		case <-pingTicker.C:
			ws.SetWriteDeadline(time.Now().Add(writeWait))
			if err := ws.WriteMessage(websocket.PingMessage, []byte{}); err != nil {
				return
			}
		}
	}
}

func (s *WebsocketServer) handleWebsocket(conn *websocket.Conn) {
	//out.Debug("New WS connection: %s", conn)
	var serializer Serializer
	var payloadType int
	if proto, ok := s.protocols[conn.Subprotocol()]; ok {
		serializer = proto.serializer
		payloadType = proto.payloadType
	} else {
		// TODO: this will not currently ever be hit because
		//       gorilla/websocket will reject the conncetion
		//       if the subprotocol isn't registered
		switch conn.Subprotocol() {
		case jsonWebsocketProtocol:
			serializer = new(JSONSerializer)
			payloadType = websocket.TextMessage
		case msgpackWebsocketProtocol:
			serializer = new(MessagePackSerializer)
			payloadType = websocket.BinaryMessage
		default:
			conn.Close()
			return
		}
	}

	peer := websocketPeer{
		conn:        conn,
		serializer:  serializer,
		messages:    make(chan Message, 10),
		payloadType: payloadType,
	}
	go peer.run()

	logErr(s.Node.Accept(&peer))
}

// Read responses from the tunnel and fulfill pending requests
func wsReader(rs *remoteServer, ws *websocket.Conn, wsTimeout time.Duration, ch chan int) {
	var err error
	log_token := cutToken(rs.token)
	// continue reading until we get an error
	for {
		ws.SetReadDeadline(time.Time{}) // no timeout, there's the ping-pong for that
		// read a message from the tunnel
		var t int
		var r io.Reader
		t, r, err = ws.NextReader()
		if err != nil {
			break
		}
		if t != websocket.BinaryMessage {
			err = fmt.Errorf("non-binary message received, type=%d", t)
			break
		}
		// give the sender a fixed time to get us the data
		ws.SetReadDeadline(time.Now().Add(wsTimeout))
		// get request id
		var id int16
		_, err = fmt.Fscanf(io.LimitReader(r, 4), "%04x", &id)
		if err != nil {
			break
		}
		// read request itself, the size is limited by the SetReadLimit on the websocket
		var buf []byte
		buf, err = ioutil.ReadAll(r)
		if err != nil {
			break
		}
		rs.log.Info("WS   RCV", "id", id, "ws", wsp(ws), "len", len(buf))
		// try to match request
		rs.requestSetMutex.Lock()
		req := rs.requestSet[id]
		rs.lastActivity = time.Now()
		rs.requestSetMutex.Unlock()
		// let's see...
		if req != nil {
			rb := responseBuffer{response: bytes.NewBuffer(buf)}
			// try to enqueue response
			select {
			case req.replyChan <- rb:
				// great!
			default:
				rs.log.Info("WS   RCV can't enqueue response", "id", id, "ws", wsp(ws))
			}
		} else {
			rs.log.Info("%s #%d: WS   RCV orphan response", "id", id, "ws", wsp(ws))
		}
	}
	// print error message
	if err != nil {
		rs.log.Info("WS   closing", "token", log_token, "err", err.Error(), "ws", wsp(ws))
	}
	// close up shop
	ch <- 0 // notify sender
	time.Sleep(2 * time.Second)
	ws.Close()
}

func newWSConn(conn *websocket.Conn, pendingWriteNum int) *WSConn {
	wsConn := new(WSConn)
	wsConn.conn = conn
	wsConn.writeChan = make(chan []byte, pendingWriteNum)

	go func() {
		for b := range wsConn.writeChan {
			if b == nil {
				break
			}

			err := conn.WriteMessage(websocket.BinaryMessage, b)
			if err != nil {
				break
			}
		}

		conn.Close()
		wsConn.Lock()
		wsConn.closeFlag = true
		wsConn.Unlock()
	}()

	return wsConn
}

func PubSubHandler(conn *websocket.Conn, pubsubClient *redis.PubSub) {
	for {
		msgi, err := pubsubClient.ReceiveMessage()

		if err != nil {
			return
		}

		switch msg := interface{}(msgi).(type) {
		case *redis.Message:
			var json_blob interface{}
			bytes_blob := []byte(msg.Payload)

			if err := json.Unmarshal(bytes_blob, &json_blob); err != nil {
				logger.Printf("[%s][error] failed to parse JSON %v, because %v", conn.RemoteAddr(), msg.Payload, err)
				continue
			}

			if err := conn.WriteJSON(json_blob); err != nil {
				logger.Printf("[%s][error] failed to send JSON, because %v", conn.RemoteAddr(), err)
				conn.Close()
				return
			}

			logger.Printf("[%s][send] OK", conn.RemoteAddr())
		default:
			logger.Printf("[%s][error] Unkown message: %s", conn.RemoteAddr(), msg)
			return
		}
	}
}

func (s *ScaletService) Create(makeFrom, rplan, name, password, location string, doStart bool,
	keys []int64, wait bool) (*Scalet, *http.Response, error) {

	scalet := new(Scalet)

	conn := new(websocket.Conn)
	var wsserr error

	if wait {
		conn, wsserr = s.client.WSSConn()
		defer conn.Close()
	}

	body := struct {
		MakeFrom string  `json:"make_from,omitempty"`
		Rplan    string  `json:"rplan,omitempty"`
		DoStart  bool    `json:"do_start,omitempty"`
		Name     string  `json:"name,omitempty"`
		Keys     []int64 `json:"keys,omitempty"`
		Password string  `json:"password,omitempty"`
		Location string  `json:"location,omitempty"`
	}{makeFrom, rplan, doStart, name, keys, password, location}

	b, _ := json.Marshal(body)

	res, err := s.client.ExecuteRequest("POST", "scalets", b, scalet)

	if wait && wsserr == nil && res.Header.Get("VSCALE-TASK-ID") != "" {
		_, err := s.client.WaitTask(conn, res.Header.Get("VSCALE-TASK-ID"))
		return scalet, res, err
	}

	return scalet, res, err
}

func writer(ws *websocket.Conn) {
	pingTicker := time.NewTicker(pingPeriod)
	nodesTicker := time.NewTicker(nodesPeriod)
	defer func() {
		pingTicker.Stop()
		nodesTicker.Stop()
		ws.Close()
	}()
	for {
		select {
		case <-nodesTicker.C:

			p := []byte(fmt.Sprintf("%v", nodesH))

			if p != nil {
				ws.SetWriteDeadline(time.Now().Add(writeWait))
				if err := ws.WriteMessage(websocket.TextMessage, p); err != nil {
					return
				}
			}
		case <-pingTicker.C:
			ws.SetWriteDeadline(time.Now().Add(writeWait))
			if err := ws.WriteMessage(websocket.PingMessage, []byte{}); err != nil {
				return
			}
		}
	}
}