siteviewcounter/main.go
Steven Polley d432521dae
All checks were successful
continuous-integration/drone/push Build is passing
wildcard CORS is bad. Make note of this, not fixed
2020-06-20 20:42:42 -06:00

133 lines
3.8 KiB
Go

/*
This application ensures the required environment variables are set
Then it uses the included countersql package to test a connection to the database
It then caches the number of unique visits in memory by querying the database
Finally, it sets up a web server with the endpoint http://localhost:8080/
For each unique source IP address which makes a GET request to this endpoint,
the count of unique visitors is incremented by 1, both in the database and in memory cache.
*/
package main
import (
"log"
"net/http"
"os"
"strconv"
"strings"
"deadbeef.codes/steven/siteviewcounter/countersql"
"github.com/go-sql-driver/mysql"
)
var (
database countersql.Configuration // Configuration information required to connect to the database, using the countersql package
uniqueVisits int // In memory cache of unique visitors
)
// Application Startup
func init() {
envVars := make(map[string]string)
envVars["dbusername"] = os.Getenv("dbusername")
envVars["dbpassword"] = os.Getenv("dbpassword")
envVars["dbhostname"] = os.Getenv("dbhostname")
envVars["dbname"] = os.Getenv("dbname")
envVars["timezone"] = os.Getenv("timezone")
for key, value := range envVars {
if value == "" {
log.Fatalf("shell environment variable %s is not set", key)
}
}
// Database Config
dbConfig := mysql.Config{}
dbConfig.User = envVars["dbusername"]
dbConfig.Passwd = envVars["dbpassword"]
dbConfig.Addr = envVars["dbhostname"]
dbConfig.DBName = envVars["dbname"]
dbConfig.Net = "tcp"
dbConfig.ParseTime = true
dbConfig.AllowNativePasswords = true
database = countersql.Configuration{}
database.DSN = dbConfig.FormatDSN()
// Test database online at startup and get count of visits
dbConn, err := database.Connect()
if err != nil {
log.Fatalf("failed to connect to database: %v", err)
}
uniqueVisits, err = dbConn.GetUniqueVisits()
if err != nil {
log.Fatalf("failed to get number of unique visits from database: %v", err)
}
dbConn.DB.Close()
}
// HTTP Routing
func main() {
// API Handlers
http.HandleFunc("/", countHandler)
log.Print("Service listening on :8080")
log.Fatal(http.ListenAndServe(":8080", nil))
}
// HTTP handler function
func countHandler(w http.ResponseWriter, r *http.Request) {
if r.Method == "GET" {
// CORS header change required.
//TBD wildcard is bad because it could allow illegitmate visits to be recorded if someone was nefarious and embedded
// front end code on a different website than your own. Need to implement environment variable to set allowed origin.
w.Header().Set("Access-Control-Allow-Origin", "*")
w.Write([]byte(strconv.Itoa(uniqueVisits)))
// Connect to database
dbConn, err := database.Connect()
if err != nil {
log.Printf("failed to connect to database: %v", err)
w.WriteHeader(http.StatusFailedDependency)
return
}
defer dbConn.DB.Close()
// We now get the source IP address of this request
var ipAddress string
// Check if we're behind a reverse proxy / WAF
if len(r.Header.Get("X-Forwarded-For")) > 0 {
ipAddress = r.Header.Get("X-Forwarded-For")
} else {
ipAddress = r.RemoteAddr
}
ipAddress = strings.Split(ipAddress, ":")[0]
// Check if this is the first time this IP address has visited
returnVisitor, err := dbConn.HasIPVisited(ipAddress)
if err != nil {
log.Printf("failed to determine if this is a return visitor, no data is being logged: %v", err)
return
}
if returnVisitor {
// Log their visit
err = dbConn.IncrementVisitor(ipAddress)
log.Printf("return visitor from %s", ipAddress)
} else {
// Insert a new visitor row in the database
err = dbConn.AddVisitor(ipAddress)
uniqueVisits++
log.Printf("new visitor from %s", ipAddress)
}
if err != nil {
log.Printf("failed to add/update visit record in database: %v", err)
return
}
} else {
// Needs to be GET method
w.WriteHeader(http.StatusMethodNotAllowed)
}
}