path: root/agent/mfi_agent.go

//
// An agent to control/report on mFi devices
//
// This agent runs on mFi devices and provides a websocket server that can be
// connected to by remote agents to control the device. The protocol is simple
// JSON over websockets. All commands will also return a report of the current
// state of the device. Reports are collected asynchronously to improve speed
// of the control channel.
//
// This daemon will also bootstrap itself into init since each reboot rewrites
// the core init files. On first run it will add itself to /etc/inittab, clean
// up some half-completed UBNT garbage that just wastes memory, and then
// return. init will start the real version of the process and keep it running.
//
// Devices only have about 5MB of free disk space and about as much free RAM so
// keeping the size of this process small is important. Go will over-allocate
// virtual memory by a factor of about 22 but the RSS of the process is
// currently about 1MB so it doesn't get OOM killed.
//
// To compile, strip and compress:
//
//   GOOS=linux GOARCH=mips go build -ldflags="-s -w" mfi_agent.go
//   upx mfi_agent
//

package main

import (
	"fmt"
	"github.com/gorilla/websocket"
	"io/ioutil"
	"log"
	"net/http"
	"os/exec"
	"path/filepath"
	"strconv"
	"time"
)

var upgrader = websocket.Upgrader{}

type Command struct {
	Cmd    string
	Args   []string
	FailOk bool
}

type PowerInfo struct {
	Output      int
	Engaged     bool
	ActivePower float64
	CurrentRMS  float64
	VoltageRMS  float64
	PowerFactor float64
	EnergySum   float64
	status      string
}

type CommandMessage struct {
	Type   string
	Output int  `json:",omitempty"`
	Engage bool `json:",omitempty"`
}

// Collect and deliver reports asynchronously. Doing these in-line casues about
// 1 second delay in the control channel.
func reporter(rchan chan chan []*PowerInfo, done <-chan bool) {
	var report []*PowerInfo

	ticker := time.NewTicker(time.Second)
	defer ticker.Stop()

	for {
		select {
		case <-ticker.C:
			report, _ = makeReport()
		case rc := <-rchan:
			rc <- report
		case <-done:
			return
		}
	}
}

// Websocket connection handler. Sends receives JSON over a websocket
// connection. Will return an error if the connection can't be upgraded.
func mfiController(rchan chan chan []*PowerInfo, w http.ResponseWriter, r *http.Request) {
	c, err := upgrader.Upgrade(w, r, nil)
	if err != nil {
		log.Println("ERROR: upgrade:", err)
		return
	}
	defer c.Close()

	res := make(chan []*PowerInfo)

	for {
		var message CommandMessage
		err = c.ReadJSON(&message)
		if err != nil {
			log.Println("ERROR: read:", err)
			break
		}

		if message.Type == "set" {
			setRelay(message.Output, message.Engage)
		}

		rchan <- res
		report := <-res

		err = c.WriteJSON(report)
		if err != nil {
			log.Println("ERROR: write:", err)
			break
		}
	}
}

// /dev/power* devices contain a report of the current sensor state of the
// PL7223 chips. This function reads them and converts one device and converts
// the output into a PowerInfo struct.
func parseOutput(path string) (*PowerInfo, error) {
	c, err := ioutil.ReadFile(path)
	if err != nil {
		return nil, err
	}

	idx, err := strconv.Atoi(string(path[len(path)-1]))
	if err != nil {
		return nil, err
	}

	out := &PowerInfo{Output: idx}
	fmt.Sscanf(
		string(c), "%s %f\n %f\n %f\n %f\n %f",
		&out.status, &out.ActivePower, &out.CurrentRMS,
		&out.VoltageRMS, &out.PowerFactor, &out.EnergySum)

	out.Engaged = (out.status == "on")

	return out, nil
}

// Gather reports for all PL7223 devices in the system
func makeReport() ([]*PowerInfo, error) {
	files, err := filepath.Glob("/dev/power?")
	if err != nil {
		return nil, err
	}

	reports := make([]*PowerInfo, 0, len(files))

	for _, fn := range files {
		o, err := parseOutput(fn)
		if err != nil {
			return nil, err
		}
		reports = append(reports, o)
	}

	return reports, nil
}

func setRelay(n int, on bool) error {
	value := []byte{'0'}
	if on {
		value[0] = '1'
	}

	return ioutil.WriteFile(
		fmt.Sprintf("/proc/power/relay%d", n), value, 0)
}

// Kill off some junk processes that don't really do anything on the current
// system to save RAM. Add ourselves to /etc/inittab and remove some more junk
// processes that inittab will continually respawn otherwise. Then HUP init to
// inform it of the config file changes. Finally kill off the remaining
// processes that we just commented out. /etc is transient and overwritten on
// every boot so we have to do this every time.
func bootstrap() error {
	cmds := []Command{
		Command{"pkill", []string{"-9", "upnpd"}, true},
		Command{"pkill", []string{"-9", "avahi"}, true},
		Command{"sh", []string{"-c", "echo null::respawn:/var/etc/persistent/mfi_agent >> /etc/inittab"}, false},
		Command{"sed", []string{"-i", "-e", "/ubnt-websockets/s/^/#/", "-e", "/telnetd/s/^/#/", "/etc/inittab"}, false},
		Command{"kill", []string{"-HUP", "1"}, false},
		Command{"pkill", []string{"-9", "ubnt-websockets"}, true},
		Command{"pkill", []string{"-9", "telnetd"}, true},
	}

	for _, cmd := range cmds {
		if err := exec.Command(cmd.Cmd, cmd.Args...).Run(); err != nil && !cmd.FailOk {
			log.Fatalf("ERROR: cmd: %s %s", cmd, err)
			return err
		}
	}

	return nil
}

// Check if we're configured in init
func isConfigured() bool {
	if err := exec.Command("grep", "mfi_agent", "/etc/inittab").Run(); err != nil {
		return false
	}
	return true
}

func main() {
	if !isConfigured() {
		bootstrap()
		return
	}

	done := make(chan bool)
	rc := make(chan chan []*PowerInfo)

	log.Printf("Running server")
	go reporter(rc, done)

	http.HandleFunc("/", func(w http.ResponseWriter, r *http.Request) {
		mfiController(rc, w, r)
	})
	log.Fatal(http.ListenAndServe("0.0.0.0:9090", nil))
	close(done)
}