Golang Receiver.Receive方法代码示例

func CmdRender(args []string, stdout, stderr io.Writer, in beam.Receiver, out beam.Sender) {
	if len(args) != 2 {
		fmt.Fprintf(stderr, "Usage: %s FORMAT\n", args[0])
		out.Send(data.Empty().Set("status", "1").Bytes(), nil)
		return
	}
	txt := args[1]
	if !strings.HasSuffix(txt, "\n") {
		txt += "\n"
	}
	t := template.Must(template.New("render").Parse(txt))
	for {
		payload, attachment, err := in.Receive()
		if err != nil {
			return
		}
		msg, err := data.Decode(string(payload))
		if err != nil {
			fmt.Fprintf(stderr, "decode error: %v\n")
		}
		if err := t.Execute(stdout, msg); err != nil {
			fmt.Fprintf(stderr, "rendering error: %v\n", err)
			out.Send(data.Empty().Set("status", "1").Bytes(), nil)
			return
		}
		if err := out.Send(payload, attachment); err != nil {
			return
		}
	}
}

func CmdDevnull(args []string, stdout, stderr io.Writer, in beam.Receiver, out beam.Sender) {
	for {
		_, attachment, err := in.Receive()
		if err != nil {
			return
		}
		if attachment != nil {
			attachment.Close()
		}
	}
}

func CmdPass(args []string, stdout, stderr io.Writer, in beam.Receiver, out beam.Sender) {
	for {
		payload, attachment, err := in.Receive()
		if err != nil {
			return
		}
		if err := out.Send(payload, attachment); err != nil {
			if attachment != nil {
				attachment.Close()
			}
			return
		}
	}
}

func CmdConnect(args []string, stdout, stderr io.Writer, in beam.Receiver, out beam.Sender) {
	if len(args) != 2 {
		out.Send(data.Empty().Set("status", "1").Set("message", "wrong number of arguments").Bytes(), nil)
		return
	}
	u, err := url.Parse(args[1])
	if err != nil {
		out.Send(data.Empty().Set("status", "1").Set("message", err.Error()).Bytes(), nil)
		return
	}
	var tasks sync.WaitGroup
	for {
		_, attachment, err := in.Receive()
		if err != nil {
			break
		}
		if attachment == nil {
			continue
		}
		Logf("connecting to %s/%s\n", u.Scheme, u.Host)
		conn, err := net.Dial(u.Scheme, u.Host)
		if err != nil {
			out.Send(data.Empty().Set("cmd", "msg", "connect error: "+err.Error()).Bytes(), nil)
			return
		}
		out.Send(data.Empty().Set("cmd", "msg", "connection established").Bytes(), nil)
		tasks.Add(1)
		go func(attachment *os.File, conn net.Conn) {
			defer tasks.Done()
			// even when successful, conn.File() returns a duplicate,
			// so we must close the original
			var iotasks sync.WaitGroup
			iotasks.Add(2)
			go func(attachment *os.File, conn net.Conn) {
				defer iotasks.Done()
				io.Copy(attachment, conn)
			}(attachment, conn)
			go func(attachment *os.File, conn net.Conn) {
				defer iotasks.Done()
				io.Copy(conn, attachment)
			}(attachment, conn)
			iotasks.Wait()
			conn.Close()
			attachment.Close()
		}(attachment, conn)
	}
	tasks.Wait()
}

func SendToConn(connections chan net.Conn, src beam.Receiver) error {
	var tasks sync.WaitGroup
	defer tasks.Wait()
	for {
		payload, attachment, err := src.Receive()
		if err == io.EOF {
			return nil
		} else if err != nil {
			return err
		}
		conn, ok := <-connections
		if !ok {
			break
		}
		Logf("Sending %s\n", msgDesc(payload, attachment))
		tasks.Add(1)
		go func(payload []byte, attachment *os.File, conn net.Conn) {
			defer tasks.Done()
			if _, err := conn.Write([]byte(data.EncodeString(string(payload)))); err != nil {
				return
			}
			if attachment == nil {
				conn.Close()
				return
			}
			var iotasks sync.WaitGroup
			iotasks.Add(2)
			go func(attachment *os.File, conn net.Conn) {
				defer iotasks.Done()
				Debugf("copying the connection to [%d]\n", attachment.Fd())
				io.Copy(attachment, conn)
				attachment.Close()
				Debugf("done copying the connection to [%d]\n", attachment.Fd())
			}(attachment, conn)
			go func(attachment *os.File, conn net.Conn) {
				defer iotasks.Done()
				Debugf("copying [%d] to the connection\n", attachment.Fd())
				io.Copy(conn, attachment)
				conn.Close()
				Debugf("done copying [%d] to the connection\n", attachment.Fd())
			}(attachment, conn)
			iotasks.Wait()
		}(payload, attachment, conn)
	}
	return nil
}

func CmdPrint(args []string, stdout, stderr io.Writer, in beam.Receiver, out beam.Sender) {
	for {
		payload, a, err := in.Receive()
		if err != nil {
			return
		}
		// Skip commands
		if a != nil && data.Message(payload).Get("cmd") == nil {
			dup, err := beam.SendRPipe(out, payload)
			if err != nil {
				a.Close()
				return
			}
			io.Copy(io.MultiWriter(os.Stdout, dup), a)
			dup.Close()
		} else {
			if err := out.Send(payload, a); err != nil {
				return
			}
		}
	}
}