Golang Channel.Write方法代码示例

// Payload: int: command size, string: command
func handleExec(ch ssh.Channel, req *ssh.Request) {
	command := string(req.Payload[4:])

	client, session, err := connectUpstream()
	if err != nil {
		ch.Write([]byte("fail to connect upstream: " + err.Error() + "\r\n"))
		ch.Close()
		return
	}

	exitStatus, err := pipe(ch, client, session, command)
	if err != nil {
		ch.Write([]byte("fail to pipe command:" + err.Error()))
		ch.Close()
		return
	}

	exitStatusBuffer := make([]byte, 4)
	binary.PutUvarint(exitStatusBuffer, uint64(exitStatus))
	log.Println("forward exit-code", exitStatus, "to client")
	_, err = ch.SendRequest("exit-status", false, exitStatusBuffer)
	if err != nil {
		log.Println("Failed to forward exit-status to client:", err)
	}

	ch.Close()
	client.Close()
	log.Println("End of exec")
}

func pipe(ch ssh.Channel, client *ssh.Client, session *ssh.Session, command string) (int, error) {
	targetStderr, err := session.StderrPipe()
	if err != nil {
		return -1, errors.New("fail to pipe stderr: " + err.Error())
	}
	targetStdout, err := session.StdoutPipe()
	if err != nil {
		return -1, errors.New("fail to pipe stdout: " + err.Error())
	}
	targetStdin, err := session.StdinPipe()
	if err != nil {
		return -1, errors.New("fail to pipe stdin: " + err.Error())
	}

	go io.Copy(targetStdin, ch)
	go io.Copy(ch.Stderr(), targetStderr)
	go io.Copy(ch, targetStdout)

	err = session.Start(command)
	if err != nil {
		ch.Write([]byte("Error when starting '" + command + "': " + err.Error()))
		ch.Close()
	}

	err = session.Wait()
	if err != nil {
		if err, ok := err.(*ssh.ExitError); ok {
			return err.ExitStatus(), nil
		} else {
			return -1, errors.New("failed to wait ssh command: " + err.Error())
		}
	}

	return 0, nil
}

func (s *scpRequest) ParseSCPRequest(channel ssh.Channel, req *ssh.Request) error {
	var err error

	// Parse the payload received from the scp client
	if s.original, err = parsePayload(req.Payload); err != nil {
		return err
	}
	log.Println("being sent file:", s.original)

	// Acknowledge payload.
	if _, err = channel.Write(zeroByte); err != nil {
		return errors.Wrap(err, "failed to write")
	}

	// Receive SCP Header
	scpHeader := make([]byte, 2048) // size of buf in openssh
	if _, err = channel.Read(scpHeader); err != nil {
		return errors.Wrap(err, "failed to retrieve header")
	}

	if _, s.size, _, err = parseHeader(scpHeader); err != nil {
		return errors.Wrap(err, "failed to parse scp header")
	}

	// Acknowledge We have received the SCP Header
	if _, err = channel.Write(zeroByte); err != nil {
		return errors.Wrap(err, "failed to reply to scp header")
	}

	return nil
}

// Ping handles a simple test SSH exec.
//
// Returns the string PONG and exit status 0.
//
// Params:
// 	- channel (ssh.Channel): The channel to respond on.
// 	- request (*ssh.Request): The request.
//
func Ping(channel ssh.Channel, req *ssh.Request) error {
	log.Info("PING")
	if _, err := channel.Write([]byte("pong")); err != nil {
		log.Err("Failed to write to channel: %s", err)
	}
	sendExitStatus(0, channel)
	req.Reply(true, nil)
	return nil
}

func sendCmdResult(channel ssh.Channel, result []byte, statusCode uint32) error {
	if _, err := channel.Write(result); err != nil {
		return fmt.Errorf("failed to write to ssh-channel: %v", err)
	}
	status := struct {
		Status uint32
	}{
		statusCode,
	}
	_, err := channel.SendRequest("exit-status", false, ssh.Marshal(&status))
	if err != nil {
		return fmt.Errorf("failed to SendRequest: %v", err)
	}
	return nil
}

// Payload: int: command size, string: command
func handleExec(ch ssh.Channel, req *ssh.Request) {
	command := string(req.Payload[4:])
	gitCmds := []string{"git-receive-pack", "git-upload-pack"}

	valid := false
	for _, cmd := range gitCmds {
		if strings.HasPrefix(command, cmd) {
			valid = true
		}
	}
	if !valid {
		ch.Write([]byte("command is not a GIT command\r\n"))
		ch.Close()
		return
	}

	ch.Write([]byte("well done!\r\n"))
	ch.Close()
}

func handleExec(ch ssh.Channel, req *ssh.Request) {
	cmd := string(req.Payload[4:])
	log.Println("received cmd", cmd)

	var msg string
	switch cmd {
	case "open", "close":
		err := sendCommand(cmd)
		if err != nil {
			log.Println("could not write to ctrl sock:", err)
			msg = "error: could not write to socket"
		} else {
			log.Printf("sent command '%s' to ctrl", cmd)
			msg = "ok"
		}
	default:
		msg = "invalid command " + cmd
	}

	ch.Write([]byte(msg + "\r\n"))
}

func (server *Server) handleRequest(channel ssh.Channel, requests <-chan *ssh.Request, conn *ssh.ServerConn) {
	defer func() {
		err := channel.Close()
		if err != nil {
			server.Logger.Errorf("Failed to close SSH Channel from %s due to %s",
				conn.RemoteAddr().String(), err)
		}
		server.Logger.Debugf("Close SSH Channel from %s", conn.RemoteAddr().String())
	}()
	for req := range requests {
		server.Logger.Debugf("Received new SSH Request (type = %s) from %s", req.Type, conn.RemoteAddr().String())

		switch req.Type {
		case "exec":
			server.handleExecRequest(channel, req, conn)
		default:
			var err error
			if req.Type == "env" {
				_, err = channel.Stderr().Write([]byte("error: Pages does not support SendEnv.\n"))
			} else {
				_, err = channel.Write([]byte("You've successfully authenticated, but Pages does not provide shell access.\n"))
			}
			if err != nil && err != io.EOF {
				server.Logger.Errorf("Failed to Talk to SSH Request due to %s", err)
			}
			err = req.Reply(false, nil)
			if err != nil && err != io.EOF {
				server.Logger.Errorf("Failed to Reply false to SSH Request due to %s", err)
			}
			err = channel.Close()
			if err != nil && err != io.EOF {
				server.Logger.Errorf("Failed to close SSH Request due to %s", err)
			}
			server.Logger.Errorf("Close SSH Request due to unsupported SSH Request type: %s", req.Type)
		}
		return
	}
}

func (s *shellHandler) startTerminal(parentTomb tomb.Tomb, sshConn *ssh.ServerConn, channel ssh.Channel) error {
	defer channel.Close()

	prompt := ">>> "
	term := terminal.NewTerminal(channel, prompt)

	// // Try to make the terminal raw
	// oldState, err := terminal.MakeRaw(0)
	// if err != nil {
	//     logger.Warn("Error making terminal raw: ", err.Error())
	// }
	// defer terminal.Restore(0, oldState)

	// Get username
	username, ok := sshConn.Permissions.Extensions["username"]
	if !ok {
		username = "user"
	}

	// Write ascii text
	term.Write([]byte(fmt.Sprintf("\r\n Nice job, %s! You are connected!\r\n", username)))
	defer term.Write([]byte(fmt.Sprintf("\r\nGoodbye, %s!\r\n", username)))

	// Start REPL
	for {

		select {
		case <-parentTomb.Dying():
			return nil
		default:
			s.logger.Info("Reading line...")
			input, err := term.ReadLine()
			if err != nil {
				fmt.Errorf("Readline() error")
				return err
			}

			// Process line
			line := strings.TrimSpace(input)
			if len(line) > 0 {

				// Log input and handle exit requests
				if line == "exit" || line == "quit" {
					s.logger.Info("Closing connection")
					return nil
				}

				// Echo input
				channel.Write(term.Escape.Green)
				channel.Write([]byte(line + "\r\n"))
				channel.Write(term.Escape.Reset)
			}
		}
	}
	return nil
}

func (e *EchoHandler) Handle(t tomb.Tomb, conn *ssh.ServerConn, channel ssh.Channel, requests <-chan *ssh.Request) error {
	defer channel.Close()
	e.logger.Info("echo handle called!")

	// Create tomb for terminal goroutines
	var tmb tomb.Tomb

	type msg struct {
		line     []byte
		isPrefix bool
		err      error
	}

	in := make(chan msg)
	defer close(in)
	reader := bufio.NewReader(channel)
	tmb.Go(func() error {
		tmb.Go(func() error {
			for {
				line, pre, err := reader.ReadLine()
				if err != nil {
					tmb.Kill(nil)
					return nil
				}

				select {
				case in <- msg{line, pre, err}:
				case <-t.Dying():
					tmb.Kill(nil)
					return nil
				case <-tmb.Dying():
					return nil
				}
			}
		})

		tmb.Go(func() error {
			for {
				e.logger.Info("time: ", time.Now())
				select {
				case <-tmb.Dying():
					return nil
				case <-t.Dying():
					tmb.Kill(nil)
					return nil
				case m := <-in:
					if m.err != nil {
						tmb.Kill(m.err)
						return m.err
					}

					// Send echo
					channel.Write(m.line)
				}
			}
		})
		return nil
	})

	return tmb.Wait()
}

// Receive receives a Git repo.
// This will only work for git-receive-pack.
func Receive(
	repo, operation, gitHome string,
	channel ssh.Channel,
	fingerprint, username, conndata, receivetype string) error {

	log.Info("receiving git repo name: %s, operation: %s, fingerprint: %s, user: %s", repo, operation, fingerprint, username)

	if receivetype == "mock" {
		channel.Write([]byte("OK"))
		return nil
	}
	repoPath := filepath.Join(gitHome, repo)
	log.Info("creating repo directory %s", repoPath)
	if _, err := createRepo(repoPath); err != nil {
		err = fmt.Errorf("Did not create new repo (%s)", err)

		return err
	}

	log.Info("writing pre-receive hook under %s", repoPath)
	if err := createPreReceiveHook(gitHome, repoPath); err != nil {
		err = fmt.Errorf("Did not write pre-receive hook (%s)", err)
		return err
	}

	cmd := exec.Command("git-shell", "-c", fmt.Sprintf("%s '%s'", operation, repo))
	log.Info(strings.Join(cmd.Args, " "))

	var errbuff bytes.Buffer

	cmd.Dir = gitHome
	cmd.Env = []string{
		fmt.Sprintf("RECEIVE_USER=%s", username),
		fmt.Sprintf("RECEIVE_REPO=%s", repo),
		fmt.Sprintf("RECEIVE_FINGERPRINT=%s", fingerprint),
		fmt.Sprintf("SSH_ORIGINAL_COMMAND=%s '%s'", operation, repo),
		fmt.Sprintf("SSH_CONNECTION=%s", conndata),
	}
	cmd.Env = append(cmd.Env, os.Environ()...)

	log.Debug("Working Dir: %s", cmd.Dir)
	log.Debug("Environment: %s", strings.Join(cmd.Env, ","))

	inpipe, err := cmd.StdinPipe()
	if err != nil {
		return err
	}
	cmd.Stdout = channel
	cmd.Stderr = io.MultiWriter(channel.Stderr(), &errbuff)

	if err := cmd.Start(); err != nil {
		err = fmt.Errorf("Failed to start git pre-receive hook: %s (%s)", err, errbuff.Bytes())
		return err
	}

	if _, err := io.Copy(inpipe, channel); err != nil {
		err = fmt.Errorf("Failed to write git objects into the git pre-receive hook (%s)", err)
		return err
	}

	fmt.Println("Waiting for git-receive to run.")
	fmt.Println("Waiting for deploy.")
	if err := cmd.Wait(); err != nil {
		err = fmt.Errorf("Failed to run git pre-receive hook: %s (%s)", errbuff.Bytes(), err)
		return err
	}
	if errbuff.Len() > 0 {
		log.Err("Unreported error: %s", errbuff.Bytes())
		return errors.New(errbuff.String())
	}
	log.Info("Deploy complete.")

	return nil
}

func HandleTcpReading(channel ssh.Channel, term *terminal.Terminal, perms *ssh.Permissions) {
	defer channel.Close()
	//http := map[string]string{}
	for {
		// read up to 1MB of data
		b := make([]byte, 1024*1024)
		_, err := channel.Read(b)
		if err != nil {
			if err.Error() == "EOF" {
				return
			}
		}
		read := bufio.NewReader(strings.NewReader(string(b)))
		toReq, err := http.ReadRequest(read)
		// TODO: https will panic atm - need to figure this out
		if err != nil {
			log.Println("Error parsing request: ", err)
			return
		}
		err = toReq.ParseForm()
		if err != nil {
			log.Println("Error parsing form: ", err)
			return
		}
		url := fmt.Sprintf("%s%s", toReq.Host, toReq.URL)

		httpReq := &HttpRequest{
			Headers:  toReq.Header,
			URL:      url,
			FormData: toReq.Form,
			Method:   toReq.Method,
			Guid:     perms.Extensions["guid"],
			Hostname: toReq.Host,
		}

		client := &http.Client{}
		resp, err := client.Get(fmt.Sprintf("http://%s", url))
		if err != nil {
			log.Fatalf("Body read error: %s", err)
		}

		defer resp.Body.Close()
		body, err2 := ioutil.ReadAll(resp.Body)
		if err2 != nil {
			log.Fatalf("Body read error: %s", err2)
		}
		httpReq.Response = string(body)
		httpReq.Save()

		log.Printf("[ http://%s ] %s", url, body)

		channel.Write(body)
		// make the http request

		//if resp, ok := httpHandler[url]; ok {
		//	channel.Write(resp)
		//} else {
		//	channel.Write([]byte("45.4.5.6"))
		//}
		channel.Close()
	}
}

func wrc(c ssh.Channel, bs []byte) error {
	_, e := c.Write(bs)
	return e
}

func (e *EchoHandler) Handle(parentTomb tomb.Tomb, sshConn *ssh.ServerConn, channel ssh.Channel, requests <-chan *ssh.Request) error {
	defer channel.Close()

	// Create tomb for terminal goroutines
	var t tomb.Tomb

	type msg struct {
		length uint32
		data   []byte
	}
	in := make(chan msg)
	defer close(in)

	// Sessions have out-of-band requests such as "shell",
	// "pty-req" and "env".  Here we handle only the
	// "shell" request.
	t.Go(func() error {
		var buffer bytes.Buffer

		// Read channel
		t.Go(func() error {

			length := make([]byte, 4)
			for {
				n, err := channel.Read(length)
				if err != nil {
					return err
				} else if n != 4 {
					return errors.New("Invalid message length")
				}

				// Decode length
				l, err := xbinary.LittleEndian.Uint32(length, 0)
				if err != nil {
					return err
				}

				// Read data
				n64, err := buffer.ReadFrom(io.LimitReader(channel, int64(l)))
				if err != nil {
					return err
				} else if n64 != int64(l) {
					return errors.New("error: reading message")
				}

				select {
				case <-parentTomb.Dying():
					return nil
				case <-t.Dying():
					return nil
				case in <- msg{l, buffer.Bytes()}:
				}
			}
		})

		length := make([]byte, 4)
	OUTER:
		for {
			select {
			case <-parentTomb.Dying():
				t.Kill(nil)
				break OUTER

			case m := <-in:
				if m.length == 0 {
					return nil
				}

				// Encode length
				_, err := xbinary.LittleEndian.PutUint32(length, 0, m.length)
				if err != nil {
					t.Kill(err)
					return nil
				}

				// Write echo response
				channel.Write(length)
				channel.Write(m.data)
			}
		}
		return nil
	})
	return t.Wait()
}

func (s *shellHandler) startTerminal(parentTomb tomb.Tomb, channel ssh.Channel, system datamodel.System, user datamodel.User) {
	defer channel.Close()

	prompt := "kappa> "
	term := terminal.NewTerminal(channel, prompt)

	// // Try to make the terminal raw
	// oldState, err := terminal.MakeRaw(0)
	// if err != nil {
	//     logger.Warn("Error making terminal raw: ", err.Error())
	// }
	// defer terminal.Restore(0, oldState)

	// Write ascii text
	term.Write([]byte("\r\n"))
	for _, line := range common.ASCII {
		term.Write([]byte(line))
		term.Write([]byte("\r\n"))
	}

	// Write login message
	term.Write([]byte("\r\n\n"))
	client.GetMessage(channel, common.DefaultColorCodes)
	term.Write([]byte("\n"))

	// Create query executor
	executor := executor.NewExecutor(executor.NewSession("", user), common.NewTerminal(term, prompt), system)

	// Start REPL
	for {

		select {
		case <-parentTomb.Dying():
			return
		default:
			input, err := term.ReadLine()
			if err != nil {
				fmt.Errorf("Readline() error")
				break
			}

			// Process line
			line := strings.TrimSpace(input)
			if len(line) > 0 {

				// Log input and handle exit requests
				if line == "exit" || line == "quit" {
					s.logger.Info("Closing connection")
					break
				} else if line == "quote me" {
					term.Write([]byte("\r\n"))
					client.GetMessage(channel, common.DefaultColorCodes)
					term.Write([]byte("\r\n"))
					continue
				} else if strings.HasPrefix(line, "//") || strings.HasPrefix(line, "--") {

					channel.Write(common.DefaultColorCodes.LightGrey)
					channel.Write([]byte(line + "\r\n"))
					channel.Write(common.DefaultColorCodes.Reset)
					continue
				}

				// Parse statement
				stmt, err := skl.ParseStatement(line)

				// Return parse error in red
				if err != nil {
					s.logger.Warn("Bad Statement", "statement", line, "error", err)
					channel.Write(common.DefaultColorCodes.LightRed)
					channel.Write([]byte(err.Error()))
					channel.Write([]byte("\r\n"))
					channel.Write(common.DefaultColorCodes.Reset)
					continue
				}

				// Execute statements
				w := common.ResponseWriter{common.DefaultColorCodes, channel}
				executor.Execute(&w, stmt)
			}
		}
	}
}