Golang unixsocket.Usock类代码示例

func sendTTYToCommand(commandUsock *unixsocket.Usock, clientFile *os.File, err error) error {
	if err != nil {
		return err
	}

	return commandUsock.WriteFD(int(clientFile.Fd()))
}

func receiveExitStatus(commandUsock *unixsocket.Usock, err error) (string, error) {
	if err != nil {
		return "", err
	}

	return commandUsock.ReadMessage()
}

func (node *SlaveNode) Run(identifier string, pid int, slaveUsock *unixsocket.Usock) {
	// TODO: We actually don't really want to prevent killing this
	// process while it's booting up.
	node.mu.Lock()
	defer node.mu.Unlock()

	node.Pid = pid

	// The slave will execute its action and respond with a status...
	msg, err := slaveUsock.ReadMessage()
	if err != nil {
		fmt.Println(err)
	}
	msg, err = ParseActionResponseMessage(msg)
	if err != nil {
		fmt.Println(err)
	}
	if msg == "OK" {
		node.Socket = slaveUsock.Conn
	} else {
		node.RegisterError(msg)
	}
	node.SignalBooted()
	slaveBooted(node.Name)

	go node.handleMessages()
}

func sendCommandLineArguments(usock *unixsocket.Usock, args []string) error {
	master, slave, err := unixsocket.Socketpair(syscall.SOCK_STREAM)
	if err != nil {
		return err
	}
	usock.WriteFD(int(slave.Fd()))
	if err != nil {
		return err
	}
	slave.Close()

	go func() {
		defer master.Close()
		argAsBytes := []byte{}
		for _, arg := range args[1:] {
			argAsBytes = append(argAsBytes, []byte(arg)...)
			argAsBytes = append(argAsBytes, byte(0))
		}
		_, err = master.Write(argAsBytes)
		if err != nil {
			slog.ErrorString("Could not send arguments across: " +
				err.Error() + "\r")
			return
		}
	}()

	return nil
}

func StartSlaveMonitor(tree *ProcessTree, local *unixsocket.Usock, remote *os.File, quit chan bool) {
	monitor := &SlaveMonitor{tree}

	// We just want this unix socket to be a channel so we can select on it...
	registeringFds := make(chan int, 3)
	go func() {
		for {
			fd, err := local.ReadFD()
			if err != nil {
				fmt.Println(err)
			}
			registeringFds <- fd
		}
	}()

	for _, slave := range monitor.tree.SlavesByName {
		if slave.Parent == nil {
			go monitor.startInitialProcess(remote)
		} else {
			go monitor.bootSlave(slave)
		}
	}

	for {
		select {
		case <-quit:
			quit <- true
			monitor.cleanupChildren()
			return
		case fd := <-registeringFds:
			monitor.slaveDidBeginRegistration(fd)
		}
	}
}

func sendExitStatus(usock *unixsocket.Usock, exitStatus string, err error) error {
	if err != nil {
		return err
	}

	_, err = usock.WriteMessage(exitStatus)
	return err
}

func receiveFileFromFD(usock *unixsocket.Usock) (*os.File, error) {
	clientFd, err := usock.ReadFD()
	if err != nil {
		return nil, errors.New("Expected FD, none received!")
	}
	fileName := strconv.Itoa(rand.Int())
	return os.NewFile(uintptr(clientFd), fileName), nil
}

func sendClientPidAndArgumentsToCommand(commandUsock *unixsocket.Usock, clientPid int, arguments string, err error) error {
	if err != nil {
		return err
	}

	msg := messages.CreatePidAndArgumentsMessage(clientPid, arguments)
	_, err = commandUsock.WriteMessage(msg)
	return err
}

func writeStacktrace(usock *unixsocket.Usock, slaveNode *processtree.SlaveNode, clientFile *os.File) {
	// Fake process ID / output / error codes:
	// Write a fake pid (step 6)
	usock.WriteMessage("0")
	// Write the error message to the terminal
	clientFile.Write([]byte(slaveNode.Error))
	// Write a non-positive exit code to the client
	usock.WriteMessage("1")
}

func sendCommandPidToClient(usock *unixsocket.Usock, pid int, err error) error {
	if err != nil {
		return err
	}

	strPid := strconv.Itoa(pid)
	_, err = usock.WriteMessage(strPid)

	return err
}

func sendClientPidAndArgumentsToCommand(commandUsock *unixsocket.Usock, clientPid int, argCount int, argFD int, err error) error {
	if err != nil {
		return err
	}

	msg := messages.CreatePidAndArgumentsMessage(clientPid, argCount)
	_, err = commandUsock.WriteMessage(msg)
	if err != nil {
		return err
	}

	return commandUsock.WriteFD(argFD)
}

func receivePidFromCommand(commandUsock *unixsocket.Usock, err error) (int, error) {
	if err != nil {
		return -1, err
	}

	msg, err := commandUsock.ReadMessage()
	if err != nil {
		return -1, err
	}
	intPid, _, _ := messages.ParsePidMessage(msg)

	return intPid, err
}

func receiveTTY(usock *unixsocket.Usock, err error) (*os.File, error) {
	if err != nil {
		return nil, err
	}

	clientFd, err := usock.ReadFD()
	if err != nil {
		return nil, errors.New("Expected FD, none received!")
	}
	fileName := strconv.Itoa(rand.Int())
	clientFile := unixsocket.FdToFile(clientFd, fileName)

	return clientFile, nil
}

func receiveCommandArgumentsAndPid(usock *unixsocket.Usock, err error) (string, int, string, error) {
	if err != nil {
		return "", -1, "", err
	}

	msg, err := usock.ReadMessage()
	if err != nil {
		return "", -1, "", err
	}
	command, clientPid, arguments, err := messages.ParseClientCommandRequestMessage(msg)
	if err != nil {
		return "", -1, "", err
	}

	return command, clientPid, arguments, err
}

// see docs/client_master_handshake.md
func handleClientConnection(tree *processtree.ProcessTree, usock *unixsocket.Usock) {
	defer usock.Close()
	// we have established first contact to the client.

	command, clientPid, argCount, argFD, err := receiveCommandArgumentsAndPid(usock, nil)
	commandNode, slaveNode, err := findCommandAndSlaveNodes(tree, command, err)
	if err != nil {
		// connection was established, no data was sent. Ignore.
		return
	}
	command = commandNode.Name // resolve aliases

	clientFile, err := receiveTTY(usock, err)
	defer clientFile.Close()

	if err == nil && slaveNode.Error != "" {
		writeStacktrace(usock, slaveNode, clientFile)
		return
	}

	commandUsock, err := bootNewCommand(slaveNode, command, err)
	if err != nil {
		// If a client connects while the command is just
		// booting up, it actually makes it here - still
		// expects a backtrace, of course.
		writeStacktrace(usock, slaveNode, clientFile)
		return
	}
	defer commandUsock.Close()

	err = sendClientPidAndArgumentsToCommand(commandUsock, clientPid, argCount, argFD, err)

	err = sendTTYToCommand(commandUsock, clientFile, err)

	cmdPid, err := receivePidFromCommand(commandUsock, err)

	err = sendCommandPidToClient(usock, cmdPid, err)

	exitStatus, err := receiveExitStatus(commandUsock, err)

	err = sendExitStatus(usock, exitStatus, err)

	if err != nil {
		slog.Error(err)
	}
	// Done! Hooray!
}