unicornpaint

A web-based painting app for raspberry PI and pimoroni Unicorn Hat HD
Log | Files | Refs | README

Server.go (4762B)


      1 // +build ignore
      2 
      3 package main
      4 
      5 import (
      6 	"encoding/json"
      7 	"fmt"
      8 
      9 	"io/ioutil"
     10 	"log"
     11 	"net/http"
     12 	"path"
     13 	"strings"
     14 
     15 	"github.com/gorilla/websocket"
     16 
     17 	"github.com/MFAshby/unicornpaint/unicorn"
     18 )
     19 
     20 var (
     21 	un unicorn.Unicorn
     22 
     23 	upgrader = websocket.Upgrader{
     24 		ReadBufferSize:  1024,
     25 		WriteBufferSize: 1024,
     26 		CheckOrigin:     func(*http.Request) bool { return true },
     27 	}
     28 	delCh       = make(chan *websocket.Conn)
     29 	addCh       = make(chan *websocket.Conn)
     30 	broadcastCh = make(chan interface{})
     31 	clients     = map[*websocket.Conn]bool{}
     32 )
     33 
     34 // Types of commands
     35 type commandType string
     36 
     37 const (
     38 	noop     commandType = "NO_OP"
     39 	setPixel commandType = "SET_PIXEL"
     40 	clear    commandType = "CLEAR"
     41 	save     commandType = "SAVE"
     42 	load     commandType = "LOAD"
     43 )
     44 
     45 type Command struct {
     46 	Type     commandType `json:"type"`
     47 	X        uint8       `json:"x"`
     48 	Y        uint8       `json:"y"`
     49 	R        uint8       `json:"r"`
     50 	G        uint8       `json:"g"`
     51 	B        uint8       `json:"b"`
     52 	SaveName string      `json:"saveName"`
     53 }
     54 
     55 const (
     56 	savesDir = "saves"
     57 )
     58 
     59 type State struct {
     60 	Saves  []string     `json:"saves"`
     61 	Pixels [][]uint8arr `json:"pixels"`
     62 }
     63 
     64 // This is a trick to avoid the JSON serializer from
     65 // interpreting uint8's as bytes and encoding them in base64
     66 type uint8arr []uint8
     67 
     68 func (u uint8arr) MarshalJSON() ([]byte, error) {
     69 	var result string
     70 	if u == nil {
     71 		result = "null"
     72 	} else {
     73 		result = strings.Join(strings.Fields(fmt.Sprintf("%d", u)), ",")
     74 	}
     75 	return []byte(result), nil
     76 }
     77 
     78 func getState() *State {
     79 	infos, err := ioutil.ReadDir(savesDir)
     80 	if err != nil {
     81 		log.Printf("Error reading saves dir %v", err)
     82 	}
     83 	saveFileNames := make([]string, len(infos))
     84 	for ix, info := range infos {
     85 		saveFileNames[ix] = info.Name()
     86 	}
     87 
     88 	// Irritating conversion function
     89 	pixels := un.GetPixels()
     90 	width := un.GetWidth()
     91 	height := un.GetHeight()
     92 
     93 	px2 := make([][]uint8arr, width)
     94 	for x := uint8(0); x < width; x++ {
     95 		px2[x] = make([]uint8arr, height)
     96 		for y := uint8(0); y < height; y++ {
     97 			px2[x][y] = uint8arr(pixels[x][y])
     98 		}
     99 	}
    100 
    101 	return &State{
    102 		Pixels: px2,
    103 		Saves:  saveFileNames,
    104 	}
    105 }
    106 
    107 func savePic(saveFileName string) {
    108 	pixels := un.GetPixels()
    109 	// Save to PNG instead
    110 	data, err := json.Marshal(pixels)
    111 	if err != nil {
    112 		log.Printf("Failed to save picture to JSON %v", err)
    113 		return
    114 	}
    115 	err = ioutil.WriteFile(path.Join(savesDir, saveFileName), data, 0644)
    116 	if err != nil {
    117 		log.Printf("Failed to write to save file %v", err)
    118 		return
    119 	}
    120 }
    121 
    122 func loadPic(saveFileName string) {
    123 	data, err := ioutil.ReadFile(path.Join(savesDir, saveFileName))
    124 	if err != nil {
    125 		log.Printf("Failed to read file %v", err)
    126 		return
    127 	}
    128 
    129 	newPixels := [][][]uint8{}
    130 	err = json.Unmarshal(data, &newPixels)
    131 	if err != nil {
    132 		log.Printf("Failed to parse file %v", err)
    133 		return
    134 	}
    135 
    136 	width := len(newPixels)
    137 	height := len(newPixels[0])
    138 	for x := 0; x < width; x++ {
    139 		for y := 0; y < height; y++ {
    140 			r, g, b := unicorn.Rgb(newPixels[x][y])
    141 			un.SetPixel(uint8(x), uint8(y), r, g, b)
    142 		}
    143 	}
    144 	un.Show()
    145 }
    146 
    147 func upgradeHandler(w http.ResponseWriter, r *http.Request) {
    148 	conn, err := upgrader.Upgrade(w, r, nil)
    149 	if err != nil {
    150 		log.Printf("Failed to upgrade someone %v", err)
    151 		return
    152 	}
    153 
    154 	// Add you to my list, remember to remove later
    155 	addCh <- conn
    156 	defer func() { delCh <- conn }()
    157 
    158 	// Get you up to speed
    159 	err = conn.WriteJSON(getState())
    160 	if err != nil {
    161 		log.Printf("Failed to send initial state %v", err)
    162 		return
    163 	}
    164 
    165 	// Read & execute commands in a loop until error
    166 	for {
    167 		cmd := Command{}
    168 		err = conn.ReadJSON(&cmd)
    169 		if err != nil {
    170 			log.Printf("Error reading from client %v", err)
    171 			break
    172 		}
    173 
    174 		switch cmd.Type {
    175 		case noop:
    176 			// Don't do anything
    177 		case setPixel:
    178 			un.SetPixel(cmd.X, cmd.Y, cmd.R, cmd.G, cmd.B)
    179 			un.Show()
    180 		case clear:
    181 			un.Clear()
    182 			un.Show()
    183 		case save:
    184 			savePic(cmd.SaveName)
    185 		case load:
    186 			loadPic(cmd.SaveName)
    187 		}
    188 		// Pretty much all commands demand a change of state,
    189 		// do a broadcast each time
    190 		broadcastCh <- getState()
    191 	}
    192 }
    193 
    194 func handleClients() {
    195 	for {
    196 		select {
    197 		case c := <-addCh:
    198 			clients[c] = true
    199 		case c := <-delCh:
    200 			delete(clients, c)
    201 		case msg := <-broadcastCh:
    202 			doBroadcast(msg)
    203 		}
    204 	}
    205 }
    206 
    207 func main() {
    208 	uni, err := unicorn.GetUnicorn()
    209 	if err != nil {
    210 		log.Fatalf("Couldn't get a unicorn :( %v", err)
    211 	}
    212 	un = uni
    213 
    214 	log.Println("Starting server on port 3001")
    215 	http.Handle("/", http.FileServer(http.Dir("build")))
    216 	http.HandleFunc("/ws", upgradeHandler)
    217 	go http.ListenAndServe(":3001", nil)
    218 	go handleClients()
    219 	un.MainLoop()
    220 }
    221 
    222 func doBroadcast(obj interface{}) {
    223 	for client := range clients {
    224 		err := client.WriteJSON(obj)
    225 
    226 		if err != nil {
    227 			log.Printf("Error writing to client, closing %v", err)
    228 			client.Close()
    229 			delete(clients, client)
    230 		}
    231 	}
    232 }