commit 64839eb22e859037a3ebc4c2fd38e65614def576 Author: Pierre-Olivier Mercier Date: Thu Nov 19 00:47:05 2020 +0100 Initial work for SRS 2020 diff --git a/db.go b/db.go new file mode 100644 index 0000000..f2e1970 --- /dev/null +++ b/db.go @@ -0,0 +1,98 @@ +package main + +import ( + "database/sql" + _ "github.com/go-sql-driver/mysql" + "log" + "os" + "time" +) + +// db stores the connection to the database +var db *sql.DB + +// DSNGenerator returns DSN filed with values from environment +func DSNGenerator() string { + db_user := "chocominer" + db_password := "chocominer" + db_host := "" + db_db := "chocominer" + + if v, exists := os.LookupEnv("MYSQL_HOST"); exists { + db_host = "tcp(" + v + ":" + if p, exists := os.LookupEnv("MYSQL_PORT"); exists { + db_host += p + ")" + } else { + db_host += "3306)" + } + } + if v, exists := os.LookupEnv("MYSQL_PASSWORD"); exists { + db_password = v + } else if v, exists := os.LookupEnv("MYSQL_ROOT_PASSWORD"); exists { + db_user = "root" + db_password = v + } + if v, exists := os.LookupEnv("MYSQL_USER"); exists { + db_user = v + } + if v, exists := os.LookupEnv("MYSQL_DATABASE"); exists { + db_db = v + } + + return db_user + ":" + db_password + "@" + db_host + "/" + db_db +} + +// DBInit establishes the connection to the database +func DBInit(dsn string) (err error) { + if db, err = sql.Open("mysql", dsn+"?parseTime=true&foreign_key_checks=1"); err != nil { + return + } + + _, err = db.Exec(`SET SESSION sql_mode = 'STRICT_TRANS_TABLES,NO_ZERO_IN_DATE,NO_ZERO_DATE,ERROR_FOR_DIVISION_BY_ZERO';`) + for i := 0; err != nil && i < 45; i += 1 { + if _, err = db.Exec(`SET SESSION sql_mode = 'STRICT_TRANS_TABLES,NO_ZERO_IN_DATE,NO_ZERO_DATE,ERROR_FOR_DIVISION_BY_ZERO';`); err != nil && i <= 45 { + log.Println("An error occurs when trying to connect to DB, will retry in 2 seconds: ", err) + time.Sleep(2 * time.Second) + } + } + + return +} + +// DBCreate creates all necessary tables used by the package +func DBCreate() error { + if _, err := db.Exec(` +CREATE TABLE IF NOT EXISTS chunks( + id_chunk INTEGER NOT NULL PRIMARY KEY AUTO_INCREMENT, + time TIMESTAMP NOT NULL, + username VARCHAR(255) NOT NULL, + chunk VARCHAR(255) NOT NULL UNIQUE, + proof VARBINARY(255) NOT NULL +) DEFAULT CHARACTER SET = utf8mb4 COLLATE = utf8mb4_unicode_ci; +`); err != nil { + return err + } + + return nil +} + +// DBClose closes the connection to the database +func DBClose() error { + return db.Close() +} + +func DBPrepare(query string) (*sql.Stmt, error) { + return db.Prepare(query) +} + +func DBQuery(query string, args ...interface{}) (*sql.Rows, error) { + return db.Query(query, args...) +} + +func DBExec(query string, args ...interface{}) (sql.Result, error) { + return db.Exec(query, args...) +} + +func DBQueryRow(query string, args ...interface{}) *sql.Row { + return db.QueryRow(query, args...) +} diff --git a/main.go b/main.go new file mode 100644 index 0000000..4750ba9 --- /dev/null +++ b/main.go @@ -0,0 +1,149 @@ +package main + +import ( + "crypto/sha256" + "encoding/hex" + "encoding/json" + "flag" + "fmt" + "log" + "math/rand" + "net/http" + "strings" + "time" +) + +var chunkSize uint = 2 +var currentChunk string + +func init() { + rand.Seed(time.Now().UnixNano()) +} + +func newChunk() { + bid := make([]byte, 4) + rand.Read(bid) + currentChunk = hex.EncodeToString(bid)[:chunkSize] +} + +type uploadChunk struct { + Proof string `json:"proof"` + Login string `json:"login"` +} + +func checkChunk(rndb string) bool { + rnd, err := hex.DecodeString(rndb) + if err != nil { + return false + } + sum := sha256.Sum256(append(rnd, currentChunk[0])) + return strings.HasPrefix(hex.EncodeToString(sum[:]), currentChunk[1:]) +} + +func ServeChunk(w http.ResponseWriter, r *http.Request) { + if addr := r.Header.Get("X-Forwarded-For"); addr != "" { + r.RemoteAddr = addr + } + log.Printf("Handling %s request from %s: %s [%s]\n", r.Method, r.RemoteAddr, r.URL.Path, r.UserAgent()) + + w.Header().Set("Content-Type", "text/plain") + + if r.Method == "POST" { + if r.ContentLength < 0 || r.ContentLength > 6553600 { + http.Error(w, fmt.Sprintf("{errmsg:\"Request too large or request size unknown\"}"), http.StatusRequestEntityTooLarge) + return + } + var body []byte + if r.ContentLength > 0 { + tmp := make([]byte, 1024) + for { + n, err := r.Body.Read(tmp) + for j := 0; j < n; j++ { + body = append(body, tmp[j]) + } + if err != nil || n <= 0 { + break + } + } + } + + var uc uploadChunk + if err := json.Unmarshal(body, &uc); err != nil { + log.Println(err) + w.WriteHeader(http.StatusBadRequest) + w.Write([]byte(err.Error())) + } else if checkChunk(uc.Proof) { + newChunk() + log.Println("Good chunk from", uc.Login) + if _, err := DBExec("INSERT INTO chunks (time, username, chunk, proof) VALUES (?, ?, ?, ?)", time.Now(), uc.Login, currentChunk, uc.Proof); err != nil { + w.WriteHeader(http.StatusBadRequest) + w.Write([]byte(err.Error())) + } else { + w.WriteHeader(http.StatusOK) + w.Write([]byte("true")) + } + } else { + log.Println("BAD chunk from", uc.Login) + } + } else { + w.WriteHeader(http.StatusOK) + w.Write([]byte(currentChunk)) + } +} + +func ServeScores(w http.ResponseWriter, r *http.Request) { + if addr := r.Header.Get("X-Forwarded-For"); addr != "" { + r.RemoteAddr = addr + } + log.Printf("Handling %s request from %s: %s [%s]\n", r.Method, r.RemoteAddr, r.URL.Path, r.UserAgent()) + + w.Header().Set("Content-Type", "text/plain") + + if rows, err := DBQuery("SELECT username, COUNT(id_chunk) AS nbchunk, MAX(time) AS time FROM chunks GROUP BY username ORDER BY nbchunk DESC"); err != nil { + w.WriteHeader(http.StatusBadRequest) + w.Write([]byte(err.Error())) + } else { + defer rows.Close() + + w.WriteHeader(http.StatusOK) + + for rows.Next() { + var login string + var nbchunk int + var time time.Time + + if err := rows.Scan(&login, &nbchunk, &time); err == nil { + w.Write([]byte(fmt.Sprintf("%q,%d,%q\n", login, nbchunk, time))) + } + } + } +} + +func main() { + var bind = flag.String("bind", "0.0.0.0:8081", "Bind port/socket") + var dsn = flag.String("dsn", DSNGenerator(), "DSN to connect to the MySQL server") + flag.UintVar(&chunkSize, "chunkSize", chunkSize, "Taille du chunk à trouver") + flag.Parse() + newChunk() + + // Database connection + log.Println("Opening database...") + if err := DBInit(*dsn); err != nil { + log.Fatal("Cannot open the database: ", err) + } + defer DBClose() + + log.Println("Creating database...") + if err := DBCreate(); err != nil { + log.Fatal("Cannot create database: ", err) + } + + log.Println("Registering handlers...") + http.HandleFunc("/scores", ServeScores) + http.HandleFunc("/chunk", ServeChunk) + + log.Println(fmt.Sprintf("Ready, listening on %s", *bind)) + if err := http.ListenAndServe(*bind, nil); err != nil { + log.Fatal("Unable to listen and serve: ", err) + } +}