siteviewcounter/main.go

133 lines
3.8 KiB
Go
Raw Normal View History

2020-06-21 02:24:51 +00:00
/*
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.
*/
2020-06-04 02:40:13 +00:00
package main
import (
"log"
"net/http"
"os"
"strconv"
"strings"
"deadbeef.codes/steven/siteviewcounter/countersql"
"github.com/go-sql-driver/mysql"
)
var (
2020-06-21 02:24:51 +00:00
database countersql.Configuration // Configuration information required to connect to the database, using the countersql package
uniqueVisits int // In memory cache of unique visitors
2020-06-04 02:40:13 +00:00
)
// 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))
}
2020-06-21 02:24:51 +00:00
// HTTP handler function
2020-06-04 02:40:13 +00:00
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.
2020-06-04 02:40:13 +00:00
w.Header().Set("Access-Control-Allow-Origin", "*")
w.Write([]byte(strconv.Itoa(uniqueVisits)))
2020-06-21 02:24:51 +00:00
// Connect to database
2020-06-04 02:40:13 +00:00
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()
2020-06-21 02:24:51 +00:00
// We now get the source IP address of this request
2020-06-04 02:40:13 +00:00
var ipAddress string
2020-06-21 02:24:51 +00:00
// Check if we're behind a reverse proxy / WAF
2020-06-04 02:40:13 +00:00
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]
2020-06-21 02:24:51 +00:00
// Check if this is the first time this IP address has visited
2020-06-04 02:40:13 +00:00
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 {
2020-06-21 02:24:51 +00:00
// Log their visit
2020-06-04 02:40:13 +00:00
err = dbConn.IncrementVisitor(ipAddress)
log.Printf("return visitor from %s", ipAddress)
} else {
2020-06-21 02:24:51 +00:00
// Insert a new visitor row in the database
2020-06-04 02:40:13 +00:00
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 {
2020-06-21 02:24:51 +00:00
// Needs to be GET method
2020-06-04 02:40:13 +00:00
w.WriteHeader(http.StatusMethodNotAllowed)
}
}