bot/bot.go

// This file is subject to a 1-clause BSD license.
// Its contents can be found in the enclosed LICENSE file.

package main

import (
	"flag"
	"log"
	"net"
	"os"
	"os/exec"
	"os/signal"
	"syscall"

	"notabug.org/mouz/bot/irc"
	"notabug.org/mouz/bot/irc/proto"
	"notabug.org/mouz/bot/plugins"

	_ "notabug.org/mouz/bot/plugins/action"
	_ "notabug.org/mouz/bot/plugins/admin"
	_ "notabug.org/mouz/bot/plugins/owm"
	_ "notabug.org/mouz/bot/plugins/knmi"
	_ "notabug.org/mouz/bot/plugins/reminder"
	_ "notabug.org/mouz/bot/plugins/url"
)

// isFork becomes true as soon as the bot has forked for the first time.
var isFork bool

// shuttingDown is true if and only if the bot is in the process of
// gracefully shutting down
var shuttingDown bool = false

func init() {
	flag.BoolVar(&isFork, "fork", false, "Whether this is a fork.")
}

// Bot defines state for a single IRC bot.
type Bot struct {
	profile irc.Profile
	client  *Client
}

// Run opens a new connection, or inherits the existing one and then begins
// the client's message poll routine.
func (b *Bot) Run() error {
	// Initialize the connection.
	err := b.Open()
	if err != nil {
		return err
	}

	// Make connection available to plugins
	irc.Connection = b.client

	// Spin up the connection's read loop.
	go func() {
		log.Println("[bot] Entering data loop...")

		err := b.client.Run()

		// err will always be non-nil here
		if e, ok := err.(*net.OpError); ok {
			if e.Err.Error() == "use of closed network connection" {
				// This can be the error value if the bot is in the
				// process of shutting down gracefully, the connection
				// is closed, and a pending read or write was
				// unblocked by that.  Just let the shutting down of
				// the bot continue and ignore the error.
				if shuttingDown {
					log.Printf("[bot] ignoring  '%+v'\n", e.Err)
					return
				}
			}
		}

		// Any other error is fatal, so a supervisor like systemd can
		// try to restart the bot.
		log.Fatal("[bot] exit 1: ", err)

	}()

	// Wait for external signals. These will either make the bot shut
	// down or start the fork.
	wait(b)
	shuttingDown = true
	return b.client.Close()
}

// PayloadHandler handles incoming server messages.
func (b *Bot) PayloadHandler(payload []byte) {
	var r irc.Request

	// Try to parse the payload into a request.
	if !parseRequest(&r, payload) {
		return
	}

	// If Target points to the bot's own name, then this message came from
	// a user as a PM. Change the Target to the sender's name, so any replies
	// we create, end up at the right destination. In any other case, the
	// target is set to the channel name from whence the message came.
	if b.profile.IsMe(r.Target) {
		r.Target = r.SenderName
	}

	// Run the appropriate handler for housekeeping.
	switch r.Type {
	case "ERROR":
		log.Println("[bot] Network error:", r.Data)
		return

	case "PING":
		proto.Pong(b.client, r.Data)
		return
	}

	// Notify plugins of message.
	plugins.Dispatch(b.client, &r)

}

// Open either establishes a new connection or inherits an existing one
// from a parent process.
func (b *Bot) Open() error {

	p := b.profile

	// Are we a fork? Then we should inherit the existing connection.
	if isFork {
		log.Println("[bot] Inherit connection to:", p.Address())

		err := b.client.OpenFd(os.NewFile(3, "conn0"))
		if err != nil {
			return err
		}

		// We're done inheriting. Have the parent process break out of
		// its wait() call by sending SIGINT to it.
		syscall.Kill(os.Getppid(), syscall.SIGINT)

		return nil
	}

	log.Println("[bot] Opening new connection to:", p.Address())

	// New bot - create a new connection.
	err := b.client.Open(p.Address())
	if err != nil {
		return err
	}

	// Perform initial handshake.
	proto.User(b.client, p.Nickname(), "8", p.Nickname())
	proto.Nick(b.client, p.Nickname(), "")
	return nil
}

// wait polls for OS signals to either kill or fork this process.
// The signals it waits for are: SIGINT, SIGTERM and SIGUSR1.
// The latter one being responsible for forking this process. The others
// are there so we may cleanly exit this process.
func wait(b *Bot) {
	signals := make(chan os.Signal, 1)
	signal.Notify(
		signals,
		syscall.SIGINT,
		syscall.SIGTERM,
		syscall.SIGUSR1,
	)

	// If the bot is just started, it should fork itself at least once
	// to play nice with systemd.  Forking is triggered by sending
	// SIGUSR1 to the current process.
	if !isFork {
		syscall.Kill(os.Getpid(), syscall.SIGUSR1)
	}

	log.Println("[bot] Waiting for signals...")

	for sig := range signals {
		log.Println("[bot] received signal:", sig)

		switch sig {

		case syscall.SIGINT, syscall.SIGTERM:
			// break out of the wait loop
			return

		case syscall.SIGUSR1:
			log.Println("[bot] forking process...")
			err := fork(b)
			if err != nil {
				log.Println("[bot]", err)
			}

		}

	}
}

// fork forks the current process into a child process and passes the
// client connection along.  The forked process is called with the
// `-fork` command line switch.
func fork(b *Bot) error {

	// Build the command line arguments for our child process.
	// This includes any custom arguments defined in the profile.
	argv := b.profile.ForkArgs()
	args := append([]string{"-fork"}, argv...)

	// Initialize the command runner.
	cmd := exec.Command(os.Args[0], args...)
	cmd.Stdout = os.Stdout
	cmd.Stderr = os.Stderr
	fd, _ := b.client.File()
	cmd.ExtraFiles = []*os.File{fd}

	// Fork the process.
	return cmd.Start()
}