Golang io.ReadWriteCloser类代码示例

// handleConnection is a per-connection go routine that registers the connection
// and starts the go routines that will read/write from the client
func handleConnection(c io.ReadWriteCloser, id string, msgchan chan m.ChatMsg, addchan chan client.Client, rmchan chan client.Client) {

	bufc := bufio.NewReader(c)
	defer c.Close()
	client := client.Client{
		Conn:     c,
		Nickname: promptNick(c, bufc),
		Ch:       make(chan m.ChatMsg),
		Id:       id,
		Kind:     "telnet",
	}

	if strings.TrimSpace(client.Nickname) == "" {
		io.WriteString(c, "Invalid Username\n")
		return
	}

	// Register user
	addchan <- client
	defer func() {
		msgchan <- m.NewChatMessage("system", "User %s left the chat room.\n", client.Nickname)
		log.Info("Connection from %v closed.\n", id)
		rmchan <- client
	}()
	io.WriteString(c, fmt.Sprintf("Welcome, %s!\n\n", client.Nickname))

	msgchan <- m.NewChatMessage("system", "New user %s has joined the chat room.\n", client.Nickname)

	// I/O
	go ReadLinesInto(client, msgchan)
	WriteLinesFrom(client)

}

func (pf *fakePortForwarder) PortForward(name string, uid types.UID, port uint16, stream io.ReadWriteCloser) error {
	defer stream.Close()

	var wg sync.WaitGroup

	// client -> server
	wg.Add(1)
	go func() {
		defer wg.Done()

		// copy from stream into a buffer
		received := new(bytes.Buffer)
		io.Copy(received, stream)

		// store the received content
		pf.lock.Lock()
		pf.received[port] = received.String()
		pf.lock.Unlock()
	}()

	// server -> client
	wg.Add(1)
	go func() {
		defer wg.Done()

		// send the hardcoded data to the stream
		io.Copy(stream, strings.NewReader(pf.send[port]))
	}()

	wg.Wait()

	return nil
}

func netcat(c io.ReadWriteCloser) {
	out("piping stdio to connection")

	done := make(chan struct{}, 2)

	go func() {
		n, _ := io.Copy(c, os.Stdin)
		out("sent %d bytes", n)
		done <- struct{}{}
	}()
	go func() {
		n, _ := io.Copy(os.Stdout, c)
		out("received %d bytes", n)
		done <- struct{}{}
	}()

	// wait until we exit.
	sigc := make(chan os.Signal, 1)
	signal.Notify(sigc, syscall.SIGHUP, syscall.SIGINT,
		syscall.SIGTERM, syscall.SIGQUIT)

	select {
	case <-done:
	case <-sigc:
		return
	}

	c.Close()
}

func setflag(c io.ReadWriteCloser) {
	defer c.Close()

	var flag_id string
	fmt.Fprintln(c, "room_id: ")
	_, err := fmt.Fscanln(c, &flag_id)
	if err != nil {
		return
	}
	var cookie string
	fmt.Fprintln(c, "auth_token: ")
	_, err = fmt.Fscanln(c, &cookie)
	if err != nil {
		return
	}
	var flag string
	fmt.Fprintln(c, "flag: ")
	_, err = fmt.Fscanln(c, &flag)
	if err != nil {
		return
	}

	if db.Set(flag_id, []string{cookie, flag}) {
		fmt.Fprintln(c, "set_flag flag_set")
		log.Println("setflag: flag set")
	} else if cookie == db.Get(flag_id)[0] {
		db.Update(flag_id, []string{cookie, flag})
		fmt.Fprintln(c, "setflag: flag_updated")
		log.Println("setflag: flag updated")
	} else {
		fmt.Fprintln(c, "setflag: flag_update_auth_fail")
		log.Println("setflag: auth fail")
	}
}

func (client *Conn) proxy(rwc io.ReadWriteCloser) (bool, error) {
	var closed bool
	for job := range client.jobs {
		if closed {
			job.results <- rillResult{err: ErrClosedConnection{}}
			continue
		}
		switch job.op {
		case _read:
			n, err := rwc.Read(job.data)
			job.results <- rillResult{n, err}
			if err != nil {
				rwc.Close()
				return false, err
			}
		case _write:
			n, err := rwc.Write(job.data)
			job.results <- rillResult{n, err}
			if err != nil {
				rwc.Close()
				return false, err
			}
		case _close:
			closed = true
			err := rwc.Close()
			job.results <- rillResult{err: err}
			return true, err
		}
	}
	return false, nil
}

// New wraps a ReadWriterCloser and uses the underlying reader and writer to
// perform a key exchange and then returns a secured connection.
func New(conn io.ReadWriteCloser) (*SecureConn, error) {
	// Generate key pair
	priv, pub, err := box.GenerateKey(rand.Reader)
	if err != nil {
		return nil, err
	}

	// Send our public key
	// This is done in a goroutine so that read/writes can happen concurrently
	// in the event that the reader and writer are directly connected (for
	// example via io.Pipe)
	done := make(chan error)
	go func() {
		_, err := conn.Write(pub[:])
		done <- err
	}()

	// Read their public key
	key := &[32]byte{}
	if _, err := io.ReadFull(conn, key[:]); err != nil {
		return nil, err
	}
	if err := <-done; err != nil {
		return nil, err
	}

	// Return a secured connection
	return &SecureConn{
		NewSecureReader(conn, priv, key),
		NewSecureWriter(conn, priv, key),
		conn,
		DefaultStreamingChunkSize,
	}, nil
}

func Pipe(src io.ReadWriteCloser, dst io.ReadWriteCloser) (int64, int64) {

	var sent, received int64
	var c = make(chan bool)
	var o sync.Once

	close := func() {
		src.Close()
		dst.Close()
		close(c)
	}

	go func() {
		received, _ = io.Copy(src, dst)
		o.Do(close)
	}()

	go func() {
		sent, _ = io.Copy(dst, src)
		o.Do(close)
	}()

	<-c
	return received, sent
}

func (srv *Server) HandleClient(raw io.ReadWriteCloser) {
	defer raw.Close()
	/*raw, err := kiss.LLObfsServerHandshake(srv.prv.PublicKey(), raw)
	if err != nil {
		kilog.Debug("directory: client failed LLObfs handshake: %v", err.Error())
		return
	}*/
	var theirKey natrium.EdDSAPublic
	// "verifier" that copies down their public key
	copier := func(haha natrium.EdDSAPublic) bool {
		theirKey = haha
		return true
	}
	sock, err := kiss.KiSSNamedHandshake(srv.prv, copier, raw)
	if err != nil {
		kilog.Debug("directory: client failed KiSS handshake: %v", err.Error())
		return
	}
	defer sock.Close()

	var request clntRequest
	err = struc.Unpack(sock, &request)
	if err != nil {
		kilog.Debug("directory: client sent malformed request: %v", err.Error())
		return
	}
	//fmt.Println(request)
	srv.rqDispatch(sock, request)
}

func (gc *Goctl) reader(c io.ReadWriteCloser) error {
	defer gc.logger.Info("Connection closed.")
	defer c.Close()

	gc.logger.Info("New connection.")
	for {
		buf, err := Read(c)
		if err != nil {
			gc.logger.Error("Error reading from connection.", "error", err)
			return err
		}

		cmd := strings.Split(string(buf), "\u0000")
		gc.logger.Debug("Got command.", "cmd", cmd)
		var resp string
		if h := gc.handlers[cmd[0]]; h != nil {
			resp = h.Run(gc, cmd[1:])
		} else {
			resp = fmt.Sprintf("ERROR: unknown command: '%s'.", cmd[0])
		}
		gc.logger.Debug("Responding.", "resp", resp)
		Write(c, []byte(resp))
	}
	/* NOTREACHED */
}

func shareEphPubKey(conn io.ReadWriteCloser, locEphPub *[32]byte) (remEphPub *[32]byte, err error) {
	var err1, err2 error
	var wg sync.WaitGroup
	wg.Add(2)

	go func() {
		defer wg.Done()
		_, err1 = conn.Write(locEphPub[:])
	}()

	go func() {
		defer wg.Done()
		remEphPub = new([32]byte)
		_, err2 = io.ReadFull(conn, remEphPub[:])
	}()

	wg.Wait()
	if err1 != nil {
		return nil, err1
	}
	if err2 != nil {
		return nil, err2
	}

	return remEphPub, nil
}

func (l *Listener) handleConnection(conn io.ReadWriteCloser, received chan<- *ReceivedMessage) {
	defer conn.Close()

	parser := SMTPParser()
	reader := bufio.NewReader(conn)
	writer := bufio.NewWriter(conn)

	session := new(Session)
	session.Start().WriteTo(writer)

	for {
		line, err := reader.ReadString('\n')
		if err != nil {
			l.Printf("error reading from client:", err)
			break
		}

		resp := session.Advance(parser(line))
		resp.WriteTo(writer)

		switch {
		case resp.IsClose():
			return
		case resp.NeedsData():
			resp, msg := session.ReadData(func() (string, error) {
				return reader.ReadString('\n')
			})
			resp.WriteTo(writer)
			if msg != nil {
				received <- msg
			}
		}
	}
}

func (o *Overlord) handleConnection(conn io.ReadWriteCloser) {
	defer conn.Close()
	decoder := codec.NewDecoder(conn, hAsocket)
	encoder := codec.NewEncoder(conn, hAsocket)
	for {
		var msg cocaine.Message
		if err := decoder.Decode(&msg); err != nil {
			log.Printf("protocol decoder error %v", err)
			return
		}
		switch msg.Session {
		case utilitySession:
			if msg.MsgType == handshake {
				log.Println("replying to heartbeat")
				encoder.Encode(msg)
			}
		default:
			o.mu.Lock()
			if ch, ok := o.sessions[msg.Session]; ok {
				select {
				case ch <- &msg:
				case <-time.After(time.Second * 2):
					log.Println("didn't send reply during 2 seconds")
				}
			} else {
				log.Printf("Invalid session %v", msg)
			}
			o.mu.Unlock()
		}
	}
}

func assertEcho(t *testing.T, rwc io.ReadWriteCloser, m ma.Multiaddr) {
	str := "test string"
	done := make(chan struct{})
	defer rwc.Close()

	go func() {
		defer close(done)

		_, err := fmt.Fprint(rwc, str)
		if err != nil {
			t.Error(err)
			return
		}

		buf := make([]byte, 256)
		n, err := rwc.Read(buf)
		if err != nil {
			t.Error(err)
			return
		}
		got := string(buf[:n])

		if got != str {
			t.Errorf("expected \"%s\", got \"%s\"", str, got)
		}
	}()

	select {
	case <-done:
	case <-time.After(time.Second):
		t.Errorf("assertEcho: %s timed out", m.String())
	}
}

func handleSocks(clnt io.ReadWriteCloser) {
	defer clnt.Close()
	destin, err := tinysocks.ReadRequest(clnt)
	if err != nil {
		log.Printf("problem while reading SOCKS5 request: %v", err.Error())
		return
	}

	log.Printf("requesting %v", destin)

	remote, err := net.Dial("tcp", destin)
	if err != nil {
		log.Printf("failed to connect to %v (%v)", destin, err.Error())
		tinysocks.CompleteRequest(0x04, clnt)
		return
	}
	tinysocks.CompleteRequest(0x00, clnt)

	log.Printf("succesfully connected to %v", destin)
	defer log.Printf("freed %v", destin)

	// forward between local and remote
	go func() {
		defer clnt.Close()
		defer remote.Close()
		io.Copy(remote, clnt)
	}()
	io.Copy(clnt, remote)
}

func (p *Proxy) accept(src io.ReadWriteCloser) {
	p.count++
	cid := p.count
	l := p.Fork("conn#%d", cid)

	l.Debugf("Open")

	if p.client.sshConn == nil {
		l.Debugf("No server connection")
		src.Close()
		return
	}

	remoteAddr := p.remote.RemoteHost + ":" + p.remote.RemotePort
	dst, err := chshare.OpenStream(p.client.sshConn, remoteAddr)
	if err != nil {
		l.Infof("Stream error: %s", err)
		src.Close()
		return
	}

	//then pipe
	s, r := chshare.Pipe(src, dst)
	l.Debugf("Close (sent %d received %d)", s, r)
}