gateway.go

package main

import (
	"Gateway-Server/model"
	"bytes"
	"context"
	"database/sql"
	"encoding/json"
	"fmt"
	_ "github.com/go-sql-driver/mysql"
	"github.com/gorilla/websocket"
	"github.com/segmentio/kafka-go"
	"io/ioutil"
	"net/http"
	"os"
	"strings"
	"sync"
)

// TODO: log to file

var genMap = newGenMap()
var DBIDMap = newIDMap()
var dbConn *sql.DB
var upgrader = websocket.Upgrader{}

func main() {

	dbConnect()    // Connect to the database
	defer func() { // TODO: handle SIGINT
		if err := dbConn.Close(); err != nil {
			fmt.Printf("Error closing db connection: %v\n", err)
		}
	}()

	http.HandleFunc("/ws", connectionHandler) // All incoming websocket connections use this

	if err := http.ListenAndServeTLS(
		":48820",
		"/etc/pki/tls/certs/gateway.crt",
		"/etc/pki/tls/private/gateway.key",
		nil); err != nil {
		fmt.Printf("Error serving TLS: %v\n", err)
	}
}

func dbConnect() {
	sqlConf, err := ioutil.ReadFile("/root/.mysql_goconf") // Read the db config file
	if err != nil {
		fmt.Printf("Error reading MySQL config file: %v\n", err)
		os.Exit(1) // If we can't connect to the database there's not much point in continuing
	}

	connInfo := strings.Fields(string(sqlConf))

	dbConn, err = sql.Open("mysql", fmt.Sprintf("%v:%v@tcp(%v)/%v",
		connInfo[0],
		connInfo[1],
		connInfo[2],
		connInfo[3]))

	if err != nil {
		fmt.Printf("Error connecting to the database: %v\n", err)
		// TODO: maybe not die here and probably handle the database going down and coming back up
		os.Exit(1)
	}
}

type socketConn struct {
	// Struct to handle our websocket connection
	conn  *websocket.Conn
	lock  sync.Mutex
	genId string
}

func connectionHandler(w http.ResponseWriter, r *http.Request) {
	conn, err := upgrader.Upgrade(w, r, nil) // Upgrade to websocket
	if err != nil {
		fmt.Printf("Error upgrading websocket: %v\n", err)
		return
	} else {
		fmt.Printf("Got connection from %v@%v\n", r.Header.Get(headerUUID), conn.RemoteAddr())
	}
	ws := &socketConn{conn: conn, lock: sync.Mutex{}}

	g := &generator{socket: ws}
	row := dbConn.QueryRow( // Figure out the database id of this generator
		"SELECT gen_id, server_id FROM generators WHERE server_id = ?",
		r.Header.Get(headerUUID))
	if err := row.Scan(&g.databaseId, &g.serverSn); err != nil { // Dump that info into our generator struct
		fmt.Printf("Error with database results: %v\n", err)
	} else {
		fmt.Printf("Got generator: %v\n", g)
	}
	ws.genId = g.serverSn
	genMap.add(g.serverSn, g)             // Add to map of SNs to generator IDs/ws connections
	DBIDMap.add(g.databaseId, g.serverSn) // Add to map of database IDs to SNs

	for { // Do forever
		var packet Packet
		err := conn.ReadJSON(&packet) // Get our next packet
		if err != nil {
			fmt.Printf("Error reading message: %v\n", err)
			return // End this connection
		} else {
			switch packet.Kind {
			case kindData:
				go ws.handleData(packet)
			case kindHeartbeat:
				go ws.handleHeartbeat(packet)
			default:
				fmt.Printf("Received unknown packet type: %v\n", packet.Kind)
			}
		}
	}
}

func (ws *socketConn) handleData(data Packet) {
	var msg map[string][]StampedReading // The message should always be a byte slice that will unmarshal into this
	if err := json.Unmarshal(data.Message, &msg); err != nil {
		fmt.Printf("Error unmarshalling data: %v\n", err)
	} else {
		fmt.Printf("got some data: %v\n", msg)
		ws.lock.Lock()
		var dataPackage = avro.NewDataPackage()
		dataPackage.Generator = int32(genMap.mapping[ws.genId].databaseId)
		// TODO: Have to get the generators state and region, probably in the same generator id map structure
		dataPackage.Organization = int32(1)
		dataPackage.State = avro.State(avro.StateCT)
		dataPackage.Region = avro.Region(avro.RegionNORTHEAST)
		for metric, stampedReadings := range msg {
			var metricData = avro.NewMetricData()
			metricData.Metric = metric
			for _, reading := range stampedReadings {
				var dataPoint = avro.NewDataPoint()
				dataPoint.Timestamp = int64(reading.Timestamp)
				var val = avro.UnionLongDouble{}
				switch v := (reading.Value).(type) {
				case int:
				case int8:
				case int16:
				case int32:
				case int64:
					val.Long = int64(v)
					val.UnionType = avro.UnionLongDoubleTypeEnumLong
					break
				case float32:
				case float64:
					val.Double = float64(v)
					val.UnionType = avro.UnionLongDoubleTypeEnumDouble
					break
				}
				dataPoint.Value = val
				metricData.Datapoints = append(metricData.Datapoints, dataPoint)
			}
			dataPackage.Data = append(dataPackage.Data, metricData)
		}
		out := new(bytes.Buffer)
		if err := dataPackage.Serialize(out); err != nil {
			fmt.Printf("Error serializing data: %v\n", err)
		}

		w := kafka.NewWriter(kafka.WriterConfig{
			Brokers:  []string{"sd5-data.engr.uconn.edu:9092"},
			Topic:    "GeneratorData",
			Balancer: &kafka.LeastBytes{},
		})

		if err := w.WriteMessages(context.Background(),
			kafka.Message{
				Key:   []byte("DataPackage"),
				Value: out.Bytes(),
			},
		); err != nil {
			fmt.Printf("Error writing to Kafka: %v\v", err)
		}

		if err := w.Close(); err != nil {
			fmt.Printf("Error closing kafka? Weird: %v\n", err)
		}

		if err := ws.conn.WriteJSON(Packet{data.ID, "response", []byte("received")}); err != nil {
			fmt.Printf("Error sending response: %v\n", err)
		} else {
			fmt.Printf("Sent response for packet %v\n", data.ID)
		}
		ws.lock.Unlock()
	}
}

func (ws *socketConn) handleHeartbeat(data Packet) {
	fmt.Printf("Heartbeat from %v\n", ws.conn.RemoteAddr())
	ws.lock.Lock()
	if err := ws.conn.WriteJSON(Packet{data.ID, "response", []byte("received")}); err != nil {
		fmt.Printf("Error sending response: %v\n", err)
	} else {
		fmt.Printf("Sent response for packet %v\n", data.ID)
	}
	ws.lock.Unlock()
}