Golang yamux.Session類代碼示例

// handle multiplex-ed connection
func handleMux(conn io.ReadWriteCloser, target string) {
	// stream multiplex
	var mux *yamux.Session
	config := &yamux.Config{
		AcceptBacklog:          256,
		EnableKeepAlive:        true,
		KeepAliveInterval:      30 * time.Second,
		ConnectionWriteTimeout: 30 * time.Second,
		MaxStreamWindowSize:    16777216,
		LogOutput:              os.Stderr,
	}
	m, err := yamux.Server(conn, config)
	if err != nil {
		log.Println(err)
		return
	}
	mux = m
	defer mux.Close()

	for {
		p1, err := mux.Accept()
		if err != nil {
			log.Println(err)
			return
		}
		p2, err := net.DialTimeout("tcp", target, 5*time.Second)
		if err != nil {
			log.Println(err)
			return
		}
		go handleClient(p1, p2)
	}
}

func handleCreateTunnel(session *yamux.Session, port int) (url string, err error) {
	// TODO: Attempt to bind to local port
	addr := fmt.Sprintf("127.0.0.1:%v", port)
	log.WithFields(logrus.Fields{
		"server": addr,
	}).Info("attempting to connect to server")
	local, err := net.Dial("tcp", addr)
	if err != nil {
		return "", errors.New("Could not establish connection to host port.")
	}
	log.WithFields(logrus.Fields{
		"server": addr,
	}).Info("connected to server")

	// Establish stream
	stream, err := session.Open()
	if err != nil {
		return "", errors.New("Could not create stream.")
	}

	// TODO: send control message
	log.WithFields(logrus.Fields{
		"server": addr,
	}).Info("sending TunnelRequest")
	req := &msg.TunnelRequest{
		Port: port,
	}
	if err := msg.WriteMsg(stream, req); err != nil {
		return "", err
	}

	// TODO: wait for ack w/ url
	log.WithFields(logrus.Fields{
		"server": addr,
	}).Info("waiting for TunnelReply")

	var res msg.TunnelReply
	if err := msg.ReadMsgInto(stream, &res); err != nil {
		return "", err
	}
	log.WithFields(logrus.Fields{
		"server": addr,
		"port":   port,
		"uri":    res.URI,
	}).Info("received TunnelReply")
	url = res.URI

	// TODO: Update global state for later control messages

	// TODO: launch goroutine to proxy data
	go proxy.Proxy(stream, local)

	return url, nil
}

func acceptStreams(logger log.Logger, session *yamux.Session, streamCh chan net.Conn) grim.TaskFunc {
	return func(ctx context.Context) {
		defer close(streamCh)
		for {
			select {
			case <-ctx.Done():
				return
			default:
				stream, err := session.Accept()
				if err != nil {
					if err != io.EOF {
						logger.Error("multiplex conn accept failed", "err", err)
					}
					return
				}
				streamCh <- stream
			}
		}
	}
}