<commit_before>package main

import (
    "github.com/jinzhu/gorm"
    _ "github.com/lib/pq"
    "github.com/gin-gonic/gin"
    _ "fmt"
    "net/http"
    _ "strings"
)

var (
    db            gorm.DB
    sqlConnection string
)

type User struct {
    Name string `json:"name"`
    Email string `json:"email"`
    Login string `json:"login"`
    Password string `json:"password"`
}

func Authentication(c *gin.Context){
    var json User
    if c.BindJSON(&json) == nil {
        result := db.Where("Login = ? and Password >= ?", json.Login, json.Password).Find(&json)
        if result.RowsAffected > 0 {
            c.JSON(http.StatusOK, gin.H{"status": result})
        } else {
            c.JSON(http.StatusNotFound, gin.H{"status": result})
        }
    }
}

func Register(c *gin.Context){
    var json User
    if c.BindJSON(&json) == nil {
        result := db.Where("Login = ? and Password >= ?", json.Login, json.Password).Find(&json)
        if result.RowsAffected == 0 {
            db.NewRecord(json)
            db.Create(&json)
            c.JSON(http.StatusOK, gin.H{"status": "Ok"})
        } else{
            c.JSON(http.StatusNotFound, gin.H{"status": "Error"})
        }

    }
}

func UsersList(c *gin.Context){
    users := []User{}
    result := db.Find(&users)
    c.JSON(http.StatusOK, gin.H{"status": result})
}

func main() {

    var err error
    sqlConnection = "dbname=auth sslmode=disable"

    db, err = gorm.Open("postgres", sqlConnection)
    db.DB()
    db.SingularTable(true)

    if err != nil {
        panic(err)
        return
    }

    router := gin.Default()
    router.Static("/", "./public")

    router.POST("/auth/login", Authentication)
    router.POST("/api/register", Register)
    router.POST("/api/users", UsersList)

    router.Run(":8080")
}<commit_msg>improvement code<commit_after>package main

import (
    "github.com/jinzhu/gorm"
    _ "github.com/lib/pq"
    "github.com/gin-gonic/gin"
    _ "fmt"
    "net/http"
    _ "strings"
)

var (
    db            gorm.DB
    sqlConnection string
    err           error
)

type User struct {
    Name     string    `json:"name"`
    Email    string    `json:"email"`
    Login    string    `json:"login"`
    Password string    `json:"password"`
}

func Authentication(c *gin.Context){
    var json User
    if c.BindJSON(&json) == nil {
        result := db.Where("Login = ? and Password >= ?", json.Login, json.Password).Find(&json)
        if result.RowsAffected > 0 {
            c.JSON(http.StatusOK, gin.H{"status": result})
        } else {
            c.JSON(http.StatusNotFound, gin.H{"status": result})
        }
    }
}

func Register(c *gin.Context){
    var json User
    if c.BindJSON(&json) == nil {
        result := db.Where("Login = ? and Password >= ?", json.Login, json.Password).Find(&json)
        if result.RowsAffected == 0 {
            db.NewRecord(json)
            db.Create(&json)
            c.JSON(http.StatusOK, gin.H{"status": "Ok"})
        } else{
            c.JSON(http.StatusNotFound, gin.H{"status": "Error"})
        }

    }
}

func UsersList(c *gin.Context){
    users := []User{}
    result := db.Find(&users)
    c.JSON(http.StatusOK, gin.H{"status": result})
}

func connection() {
    sqlConnection = "dbname=auth sslmode=disable"

    db, err = gorm.Open("postgres", sqlConnection)
    db.DB()
    db.SingularTable(true)

    if err != nil {
        panic(err)
        return
    }
}

func main() {
    //create connection db
    connection()

    router := gin.Default()
    router.Static("/", "./public")

    router.POST("/auth/login", Authentication)
    router.POST("/api/register", Register)
    router.POST("/api/users", UsersList)

    router.Run(":8080")
}
<|endoftext|>


<commit_before>// vim: tabstop=2 shiftwidth=2

package main

import (
	"fmt"
	"os"
	"path"
	"bytes"
	"time"
	"errors"
	"io/ioutil"
	"encoding/hex"
	"crypto/sha256"
	"crypto/sha512"
	"github.com/crooks/yamn/keymgr"
	"github.com/crooks/yamn/idlog"
	//"github.com/codahale/blake2"
)

func loopServer() (err error) {
	var filenames []string
	// Populate public and secret keyrings
	public := keymgr.NewPubring(cfg.Files.Pubring, cfg.Files.Mlist2)
	secret := keymgr.NewSecring(cfg.Files.Secring, cfg.Files.Pubkey)
	public.ImportPubring()
	secret.ImportSecring()
	// Create some dirs if it doesn't already exist
	err = os.MkdirAll(cfg.Files.IDlog, 0700)
	if err != nil {
		return
	}
	err = os.MkdirAll(cfg.Files.Pooldir, 0700)
	if err != nil {
		return
	}
	// Open the IDlog
	Trace.Printf("Opening ID Log: %s", cfg.Files.IDlog)
	idlog, err := idlog.NewInstance(cfg.Files.IDlog)
	if err != nil {
		panic(err)
	}
	defer idlog.Close()
	// Is a new ECC Keypair required?
	generate := false
	// Find out the Keyid we're advertising
	err = secret.SetMyKeyid()
	if err != nil {
		Info.Println("No valid Public key, will generate a new pair")
		generate = true
	}
	// Try to validate the advertised key on Secring
	if ! generate {
		valid, err := secret.Validate()
		if err != nil {
			Warn.Printf("%s: Failed to validate key in Secring", cfg.Files.Secring)
			generate = true
		} else if valid {
			Info.Printf("Advertising keyid=%s", secret.GetMyKeyidStr())
		} else {
			Info.Printf("%s has expired, will generate a new key pair", secret.GetMyKeyidStr())
			generate = true
		}
	}
	// Create a new key pair
	if generate {
		Info.Println("Generating and advertising a new key pair")
		pub, sec := eccGenerate()
		err = secret.Publish(
			pub, sec,
			keyValidityDays,
			cfg.Remailer.Exit,
			cfg.Remailer.Name, cfg.Remailer.Address, version)
		if err != nil {
			Error.Printf("Aborting! Key generation failure: %s", err)
			return
		}
		Info.Printf("Advertising new Keyid: %s", secret.GetMyKeyidStr())
	}

	//TODO Make this a flag function
	secret.Purge("test.txt")

	// Complain about excessively small loop values.
	if cfg.Remailer.Loop < 60 {
		Warn.Println(
			fmt.Sprintf("Loop time of %d is excessively low. ", cfg.Remailer.Loop),
			"Will loop every 60 seconds. A higher setting is recommended.")
	}
	// Complain about high pool rates.
	if cfg.Pool.Rate > 90 {
		Warn.Println(
			fmt.Sprintf("Your pool rate of %d is excessively", cfg.Pool.Rate),
			"high. Unless testing, a lower setting is recommended.")
	}

	// Maintain time of last pool process
	poolProcessTime := time.Now()
	poolProcessDelay := time.Duration(cfg.Remailer.Loop) * time.Second

	// Actually start the server loop
	Info.Printf("Starting YAMN server: %s", cfg.Remailer.Name)
	for {
		if time.Now().Before(poolProcessTime) {
			mailRead()
			// Don't do anything beyond this point until poolProcessTime
			time.Sleep(60 * time.Second)
			continue
		}
		// Test if in-memory pubring is current
		if public.KeyRefresh() {
			// Time to re-read the pubring file
			Info.Printf("Reimporting keyring: %s", cfg.Files.Pubring)
			public.ImportPubring()
		}
		filenames, err = poolRead()
		for _, file := range filenames {
			err = processPoolFile(file, secret, idlog)
			if err != nil {
				Info.Printf("Discarding message: %s", err)
			}
			poolDelete(file)
		}
		poolProcessTime = time.Now().Add(poolProcessDelay)
	}
}

func exportMessage(headers, fake, body []byte, sendto string) (err error) {
	hlen := len(headers) + len(fake)
	err = lenCheck(hlen, headersBytes)
	if err != nil {
		return
	}
	err = lenCheck(len(body), bodyBytes)
	if err != nil {
		return
	}
	buf := new(bytes.Buffer)
	buf.Write(headers)
	buf.Write(fake)
	buf.Write(body)
	err = bufLenCheck(buf.Len(), messageBytes)
	if err != nil {
		Error.Println("Incorrect outbound message size. Not sending.")
		return
	}
	if sendto == cfg.Remailer.Address {
		Info.Println("Message loops back to us. Storing in pool.")
		//digest := blake2.New(&blake2.Config{Size: 16})
		digest := sha256.New()
		digest.Write(buf.Bytes())
		filename := "m" + hex.EncodeToString(digest.Sum(nil))
		filename = path.Join(cfg.Files.Pooldir, filename[:14])
		err = ioutil.WriteFile(filename, buf.Bytes(), 0600)
		if err != nil {
			Warn.Println(err)
			return
		}
	} else {
		Trace.Printf("Forwarding message to: %s", sendto)
		err = cutmarks(buf.Bytes(), sendto)
	}
	return
}

func processPoolFile(filename string, secret *keymgr.Secring, idlog idlog.IDLog) (err error) {
	f, err := os.Open(path.Join(cfg.Files.Pooldir, filename))
	defer f.Close()
	if err != nil {
		return
	}
	Trace.Printf("Processing pool file: %s\n", filename)
	// Initialize some slices for the message components
	header := make([]byte, headerBytes)
	headers := make([]byte, encHeadBytes)
	body := make([]byte, bodyBytes)
	var bytesRead int
	// Read each message component and validate its size
	bytesRead, err = f.Read(header)
	if err != nil {
		return
	}
	if bytesRead != headerBytes {
		Warn.Printf("Incorrect header bytes. Wanted=%d, Got=%d", headerBytes, bytesRead)
		return
	}
	bytesRead, err = f.Read(headers)
	if err != nil {
		return
	}
	if bytesRead != encHeadBytes {
		Warn.Printf("Incorrect headers bytes. Wanted=%d, Got=%d", encHeadBytes, bytesRead)
		return
	}
	bytesRead, err = f.Read(body)
	if err != nil {
		return
	}
	if bytesRead != bodyBytes {
		Warn.Printf("Incorrect body bytes. Wanted=%d, Got=%d", bodyBytes, bytesRead)
		return
	}

	var iv []byte
	/*
	decodeHead only returns the decrypted slotData bytes.  The other fields are
	only concerned with performing the decryption.
	*/
	var decodedHeader []byte
	decodedHeader, err = decodeHead(header, secret)
	if err != nil {
		return
	}
	// data contains the slotData struct
	data := new(slotData)
	err = data.decodeData(decodedHeader)
	if err != nil {
		return
	}
	// Test uniqueness of packet ID
	if ! idlog.Unique(data.packetID, cfg.Remailer.IDexp) {
		err = errors.New("Packet ID collision")
		return
	}
	//digest := blake2.New(nil)
	digest := sha512.New()
	digest.Write(headers)
	digest.Write(body)
	if ! bytes.Equal(digest.Sum(nil), data.tagHash) {
		err = fmt.Errorf("Digest mismatch on Anti-tag hash")
		return
	}
	if data.packetType == 0 {
		Trace.Println("This is an Intermediate type message")
		// inter contains the slotIntermediate struct
		inter := new(slotIntermediate)
		err = inter.decodeIntermediate(data.packetInfo)
		if err != nil {
			return
		}
		// Number of headers to decrypt is one less than max chain length
		for headNum := 0; headNum < maxChainLength - 1; headNum++ {
			iv, err = sPopBytes(&inter.aesIVs, 16)
			if err != nil {
				return
			}
			sbyte := headNum * headerBytes
			ebyte := (headNum + 1) * headerBytes
			copy(headers[sbyte:ebyte], AES_CTR(headers[sbyte:ebyte], data.aesKey, iv))
		}
		// The tenth IV is used to encrypt the deterministic header
		iv, err = sPopBytes(&inter.aesIVs, 16)
		if err != nil {
			return
		}
		//fmt.Printf("Fake: Key=%x, IV=%x\n", data.aesKey[:10], iv[:10])
		fakeHeader := make([]byte, headerBytes)
		copy(fakeHeader, AES_CTR(fakeHeader, data.aesKey, iv))
		// Body is decrypted with the final IV
		iv, err = sPopBytes(&inter.aesIVs, 16)
		if err != nil {
			return
		}
		copy(body, AES_CTR(body, data.aesKey, iv))
		// At this point there should be zero bytes left in the inter IV pool
		if len(inter.aesIVs) != 0 {
			err = fmt.Errorf("IV pool not empty.  Contains %d bytes.", len(inter.aesIVs))
			return
		}
		err = exportMessage(headers, fakeHeader, body, inter.nextHop)
		if err != nil {
			return
		}
	} else if data.packetType == 1 {
		Trace.Println("This is an Exit type message")
		final := new(slotFinal)
		err = final.decodeFinal(data.packetInfo)
		if err != nil {
			return
		}
		body = AES_CTR(body[:final.bodyBytes], data.aesKey, final.aesIV)
		var recipients []string
		recipients, err = testMail(body)
		if err != nil {
			return
		}
		for _, sendto := range recipients {
			if cfg.Mail.Sendmail {
				err = sendmail(body, sendto)
				if err != nil {
					Warn.Println("Sendmail failed")
					return
				}
			} else {
				err = SMTPRelay(body, sendto)
				if err != nil {
					Warn.Println("SMTP relay failed")
					return
				}
			}
		}
	}
	return
}
<commit_msg>Make key expiration and generation a daily task<commit_after>// vim: tabstop=2 shiftwidth=2

package main

import (
	"fmt"
	"os"
	"path"
	"bytes"
	"time"
	"errors"
	"io/ioutil"
	"encoding/hex"
	"crypto/sha256"
	"crypto/sha512"
	"github.com/crooks/yamn/keymgr"
	"github.com/crooks/yamn/idlog"
	//"github.com/codahale/blake2"
)

func loopServer() (err error) {
	var filenames []string
	// Populate public and secret keyrings
	public := keymgr.NewPubring(cfg.Files.Pubring, cfg.Files.Mlist2)
	secret := keymgr.NewSecring(cfg.Files.Secring, cfg.Files.Pubkey)
	public.ImportPubring()
	secret.ImportSecring()
	// Create some dirs if it doesn't already exist
	err = os.MkdirAll(cfg.Files.IDlog, 0700)
	if err != nil {
		return
	}
	err = os.MkdirAll(cfg.Files.Pooldir, 0700)
	if err != nil {
		return
	}
	// Open the IDlog
	Trace.Printf("Opening ID Log: %s", cfg.Files.IDlog)
	idlog, err := idlog.NewInstance(cfg.Files.IDlog)
	if err != nil {
		panic(err)
	}
	defer idlog.Close()
	// Is a new ECC Keypair required?
	generate := false
	/*
	Ascertain the Keyid we're advertising.  This only needs to be done once at
	server startup as the act of new key publication sets the keyid to the newly
	generated key.
	*/
	err = secret.SetMyKeyid()
	if err != nil {
		Info.Println("No valid Public key, will generate a new pair")
		generate = true
	} else {
		Info.Printf("Advertising existing keyid: %s", secret.GetMyKeyidStr())
	}

	// Maintain time of last pool process
	poolProcessTime := time.Now()
	poolProcessDelay := time.Duration(cfg.Remailer.Loop) * time.Second

	// Make a note of what day it is
	today := time.Now()
	oneday := time.Duration(24) * time.Hour

	// Actually start the server loop
	Info.Printf("Starting YAMN server: %s", cfg.Remailer.Name)
	for {
		if time.Now().Before(poolProcessTime) {
			mailRead()
			// Don't do anything beyond this point until poolProcessTime
			time.Sleep(60 * time.Second)
			continue
		}
		// Test if it's time to do daily events
		if generate || time.Since(today) > oneday {
			Info.Println("Performing daily events")
			// Try to validate the advertised key on Secring
			if ! generate {
				valid, err := secret.Validate()
				if err != nil {
					Warn.Printf("%s: Failed to validate key in Secring", cfg.Files.Secring)
					generate = true
				} else if valid {
					Info.Printf("Advertising current keyid: %s", secret.GetMyKeyidStr())
				} else {
					Info.Printf("%s has expired, will generate a new key pair", secret.GetMyKeyidStr())
					generate = true
				}
			}
			// Create a new key pair
			if generate {
				Info.Println("Generating and advertising a new key pair")
				pub, sec := eccGenerate()
				err = secret.Publish(
					pub, sec,
					keyValidityDays,
					cfg.Remailer.Exit,
					cfg.Remailer.Name, cfg.Remailer.Address, version)
				if err != nil {
					Error.Printf("Aborting! Key generation failure: %s", err)
					return
				}
				Info.Printf("Advertising new Keyid: %s", secret.GetMyKeyidStr())
				generate = false
			}

			//TODO Make this a flag function
			secret.Purge("test.txt")

			// Complain about excessively small loop values.
			if cfg.Remailer.Loop < 60 {
				Warn.Println(
					fmt.Sprintf("Loop time of %d is excessively low. ", cfg.Remailer.Loop),
					"Will loop every 60 seconds. A higher setting is recommended.")
			}
			// Complain about high pool rates.
			if cfg.Pool.Rate > 90 {
				Warn.Println(
					fmt.Sprintf("Your pool rate of %d is excessively", cfg.Pool.Rate),
					"high. Unless testing, a lower setting is recommended.")
			}
			// Reset today so we don't do these tasks for the next 24 hours.
			today = time.Now()
		}
		// Test if in-memory pubring is current
		if public.KeyRefresh() {
			// Time to re-read the pubring file
			Info.Printf("Reimporting keyring: %s", cfg.Files.Pubring)
			public.ImportPubring()
		}
		filenames, err = poolRead()
		for _, file := range filenames {
			err = processPoolFile(file, secret, idlog)
			if err != nil {
				Info.Printf("Discarding message: %s", err)
			}
			poolDelete(file)
		}
		poolProcessTime = time.Now().Add(poolProcessDelay)
	}
}

func exportMessage(headers, fake, body []byte, sendto string) (err error) {
	hlen := len(headers) + len(fake)
	err = lenCheck(hlen, headersBytes)
	if err != nil {
		return
	}
	err = lenCheck(len(body), bodyBytes)
	if err != nil {
		return
	}
	buf := new(bytes.Buffer)
	buf.Write(headers)
	buf.Write(fake)
	buf.Write(body)
	err = bufLenCheck(buf.Len(), messageBytes)
	if err != nil {
		Error.Println("Incorrect outbound message size. Not sending.")
		return
	}
	if sendto == cfg.Remailer.Address {
		Info.Println("Message loops back to us. Storing in pool.")
		//digest := blake2.New(&blake2.Config{Size: 16})
		digest := sha256.New()
		digest.Write(buf.Bytes())
		filename := "m" + hex.EncodeToString(digest.Sum(nil))
		filename = path.Join(cfg.Files.Pooldir, filename[:14])
		err = ioutil.WriteFile(filename, buf.Bytes(), 0600)
		if err != nil {
			Warn.Println(err)
			return
		}
	} else {
		Trace.Printf("Forwarding message to: %s", sendto)
		err = cutmarks(buf.Bytes(), sendto)
	}
	return
}

func processPoolFile(filename string, secret *keymgr.Secring, idlog idlog.IDLog) (err error) {
	f, err := os.Open(path.Join(cfg.Files.Pooldir, filename))
	defer f.Close()
	if err != nil {
		return
	}
	Trace.Printf("Processing pool file: %s\n", filename)
	// Initialize some slices for the message components
	header := make([]byte, headerBytes)
	headers := make([]byte, encHeadBytes)
	body := make([]byte, bodyBytes)
	var bytesRead int
	// Read each message component and validate its size
	bytesRead, err = f.Read(header)
	if err != nil {
		return
	}
	if bytesRead != headerBytes {
		Warn.Printf("Incorrect header bytes. Wanted=%d, Got=%d", headerBytes, bytesRead)
		return
	}
	bytesRead, err = f.Read(headers)
	if err != nil {
		return
	}
	if bytesRead != encHeadBytes {
		Warn.Printf("Incorrect headers bytes. Wanted=%d, Got=%d", encHeadBytes, bytesRead)
		return
	}
	bytesRead, err = f.Read(body)
	if err != nil {
		return
	}
	if bytesRead != bodyBytes {
		Warn.Printf("Incorrect body bytes. Wanted=%d, Got=%d", bodyBytes, bytesRead)
		return
	}

	var iv []byte
	/*
	decodeHead only returns the decrypted slotData bytes.  The other fields are
	only concerned with performing the decryption.
	*/
	var decodedHeader []byte
	decodedHeader, err = decodeHead(header, secret)
	if err != nil {
		return
	}
	// data contains the slotData struct
	data := new(slotData)
	err = data.decodeData(decodedHeader)
	if err != nil {
		return
	}
	// Test uniqueness of packet ID
	if ! idlog.Unique(data.packetID, cfg.Remailer.IDexp) {
		err = errors.New("Packet ID collision")
		return
	}
	//digest := blake2.New(nil)
	digest := sha512.New()
	digest.Write(headers)
	digest.Write(body)
	if ! bytes.Equal(digest.Sum(nil), data.tagHash) {
		err = fmt.Errorf("Digest mismatch on Anti-tag hash")
		return
	}
	if data.packetType == 0 {
		Trace.Println("This is an Intermediate type message")
		// inter contains the slotIntermediate struct
		inter := new(slotIntermediate)
		err = inter.decodeIntermediate(data.packetInfo)
		if err != nil {
			return
		}
		// Number of headers to decrypt is one less than max chain length
		for headNum := 0; headNum < maxChainLength - 1; headNum++ {
			iv, err = sPopBytes(&inter.aesIVs, 16)
			if err != nil {
				return
			}
			sbyte := headNum * headerBytes
			ebyte := (headNum + 1) * headerBytes
			copy(headers[sbyte:ebyte], AES_CTR(headers[sbyte:ebyte], data.aesKey, iv))
		}
		// The tenth IV is used to encrypt the deterministic header
		iv, err = sPopBytes(&inter.aesIVs, 16)
		if err != nil {
			return
		}
		//fmt.Printf("Fake: Key=%x, IV=%x\n", data.aesKey[:10], iv[:10])
		fakeHeader := make([]byte, headerBytes)
		copy(fakeHeader, AES_CTR(fakeHeader, data.aesKey, iv))
		// Body is decrypted with the final IV
		iv, err = sPopBytes(&inter.aesIVs, 16)
		if err != nil {
			return
		}
		copy(body, AES_CTR(body, data.aesKey, iv))
		// At this point there should be zero bytes left in the inter IV pool
		if len(inter.aesIVs) != 0 {
			err = fmt.Errorf("IV pool not empty.  Contains %d bytes.", len(inter.aesIVs))
			return
		}
		err = exportMessage(headers, fakeHeader, body, inter.nextHop)
		if err != nil {
			return
		}
	} else if data.packetType == 1 {
		Trace.Println("This is an Exit type message")
		final := new(slotFinal)
		err = final.decodeFinal(data.packetInfo)
		if err != nil {
			return
		}
		body = AES_CTR(body[:final.bodyBytes], data.aesKey, final.aesIV)
		var recipients []string
		recipients, err = testMail(body)
		if err != nil {
			return
		}
		for _, sendto := range recipients {
			if cfg.Mail.Sendmail {
				err = sendmail(body, sendto)
				if err != nil {
					Warn.Println("Sendmail failed")
					return
				}
			} else {
				err = SMTPRelay(body, sendto)
				if err != nil {
					Warn.Println("SMTP relay failed")
					return
				}
			}
		}
	}
	return
}
<|endoftext|>


<commit_before>package main

import (
  "github.com/codegangsta/martini"
  "encoding/json"
  "net/http"
  "os"
  "github.com/jinzhu/gorm"
  _ "github.com/go-sql-driver/mysql"
)

func main() {
  m := martini.Classic()

  dbPassword := os.Getenv("OOPS_DB_PASS")
  db, err := gorm.Open("mysql", "oops:" + dbPassword + "@tcp(15.126.247.23:3306)/oops")
  if err != nil { panic(err) }
  //defer db.Close() //TODO what should this be?

  //db.SetPool(10)

  db.LogMode(true)

  m.Map(db)

  m.Get("/schools", func(res http.ResponseWriter) string {
    var allSchools = []School{}
  	db.Find(&allSchools)
    return render(res, allSchools)
  })

//TODO need to also return stats by year -- needs to return []DistrictsByYear - district level stats for each year and each district
  m.Get("/districts", func(res http.ResponseWriter) string {
    var allDistricts = []District{}
    db.Find(&allDistricts)
    return render(res, allDistricts)
  })

//TODO need to also return stats by year for each school in the district - needs to return SchoolsByYear[] - data for all schools in this district
  m.Get("/district/:id", func(res http.ResponseWriter, params martini.Params) string {
    var testDistrict = District{}
  	db.Where("id = ?", params["id"]).First(&testDistrict)
    return render(res, testDistrict)

  })

//TODO Possible clean-up opportunities here
//TODO Still need total teacher and student count per year - student count could be calculated from grade counts
  m.Get("/school/:id", func(res http.ResponseWriter, params martini.Params) string {
    var testSchool = School{}
  	db.Where("id = ?",params["id"]).First(&testSchool)
  	var testStats = []ClassStat{}
  	db.Where("school_id = ?", params["id"]).Find(&testStats)

    var yearsToEnrollments = map[string][]GradeEnrollment{}
  	for _, row := range testStats {
  		yearsToEnrollments[row.Years] = append(yearsToEnrollments[row.Years], GradeEnrollment{
  			Grade: row.Grade,
  			EnrollmentSize:    row.EnrollmentSize,
  		})
  	}

  	var enrollmentData = []EnrollmentByYear{}
  	for year, enrollment := range yearsToEnrollments{
  		enrollmentData = append(enrollmentData, EnrollmentByYear{
  			Year: year,
  			GradeEnrollment: enrollment,
  		})
  	}

  	var allData = SchoolWithEnrollment {
  		School: testSchool,
  		EnrollmentByYear: enrollmentData,
  	}
    return render(res, allData)
  })

  m.Run()
}

func render(res http.ResponseWriter, data interface{}) string {
  thing, err := json.Marshal(data)
  if err != nil { panic(err) }
  return asJson(res, thing)
}

func asJson(res http.ResponseWriter, data []byte) string {
  res.Header().Set("Content-Type", "application/json")
  res.Header().Set("Access-Control-Allow-Origin", "*")
  return string(data[:])
}

<commit_msg>tabs2spaces<commit_after>package main

import (
  "github.com/codegangsta/martini"
  "encoding/json"
  "net/http"
  "os"
  "github.com/jinzhu/gorm"
  _ "github.com/go-sql-driver/mysql"
)

func main() {
  m := martini.Classic()

  dbPassword := os.Getenv("OOPS_DB_PASS")
  db, err := gorm.Open("mysql", "oops:" + dbPassword + "@tcp(15.126.247.23:3306)/oops")
  if err != nil { panic(err) }
  //defer db.Close() //TODO what should this be?

  //db.SetPool(10)

  db.LogMode(true)

  m.Map(db)

  m.Get("/schools", func(res http.ResponseWriter) string {
    var allSchools = []School{}
    db.Find(&allSchools)
    return render(res, allSchools)
  })

  //TODO need to also return stats by year -- needs to return []DistrictsByYear - district level stats for each year and each district
  m.Get("/districts", func(res http.ResponseWriter) string {
    var allDistricts = []District{}
    db.Find(&allDistricts)
    return render(res, allDistricts)
  })

  //TODO need to also return stats by year for each school in the district - needs to return SchoolsByYear[] - data for all schools in this district
  m.Get("/district/:id", func(res http.ResponseWriter, params martini.Params) string {
    var testDistrict = District{}
    db.Where("id = ?", params["id"]).First(&testDistrict)
    return render(res, testDistrict)

  })

  //TODO Possible clean-up opportunities here
  //TODO Still need total teacher and student count per year - student count could be calculated from grade counts
  m.Get("/school/:id", func(res http.ResponseWriter, params martini.Params) string {
    var testSchool = School{}
    db.Where("id = ?",params["id"]).First(&testSchool)
    var testStats = []ClassStat{}
    db.Where("school_id = ?", params["id"]).Find(&testStats)

    var yearsToEnrollments = map[string][]GradeEnrollment{}
    for _, row := range testStats {
      yearsToEnrollments[row.Years] = append(yearsToEnrollments[row.Years], GradeEnrollment{
        Grade: row.Grade,
        EnrollmentSize:    row.EnrollmentSize,
      })
    }

    var enrollmentData = []EnrollmentByYear{}
    for year, enrollment := range yearsToEnrollments{
      enrollmentData = append(enrollmentData, EnrollmentByYear{
        Year: year,
        GradeEnrollment: enrollment,
      })
    }

    var allData = SchoolWithEnrollment {
      School: testSchool,
      EnrollmentByYear: enrollmentData,
    }
  return render(res, allData)
  })

  m.Run()
}

func render(res http.ResponseWriter, data interface{}) string {
  thing, err := json.Marshal(data)
  if err != nil { panic(err) }
  return asJson(res, thing)
}

func asJson(res http.ResponseWriter, data []byte) string {
  res.Header().Set("Content-Type", "application/json")
  res.Header().Set("Access-Control-Allow-Origin", "*")
  return string(data[:])
}

<|endoftext|>


<commit_before>/*
Copyright 2017 Eric Chlebek

Permission is hereby granted, free of charge, to any person obtaining a copy of
this software and associated documentation files (the "Software"), to deal in
the Software without restriction, including without limitation the rights to
use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of
the Software, and to permit persons to whom the Software is furnished to do so,
subject to the following conditions:

The above copyright notice and this permission notice shall be included in all
copies or substantial portions of the Software.

THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS
FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR
COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER
IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN
CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
*/

package main

import (
	"bytes"
	"encoding/base64"
	"encoding/json"
	"errors"
	"flag"
	"fmt"
	"hash/fnv"
	"io"
	"log"
	"net/http"

	"github.com/boltdb/bolt"
)

var (
	DBName       = flag.String("db", "bolt.db", "Bolt database to use")
	headerBucket = append([]byte{0}, []byte("headers")...)
	headerKeys   = []string{
		"Content-Type",
		"Content-Length",
		"ETag",
	}
)

type context struct {
	db *bolt.DB
}

type router struct {
	ctx context
}

func logRequest(req *http.Request) {
	log.Println(req.Method, req.URL.Path)
}

func createHeaderBucketIfNotExists(db *bolt.DB) error {
	return db.Update(func(tx *bolt.Tx) error {
		bucket, err := tx.CreateBucketIfNotExists(headerBucket)
		if err != nil {
			return err
		}
		return bucket.Put([]byte("/"), []byte("{}"))
	})
}

func createRootBucketIfNotExists(db *bolt.DB) error {
	return db.Update(func(tx *bolt.Tx) error {
		_, err := tx.CreateBucketIfNotExists([]byte("/"))
		return err
	})
}

func getBoltBucketOrValue(bucket *bolt.Bucket, key []byte) (*bolt.Bucket, []byte) {
	b := bucket.Bucket(key)
	if b == nil {
		return nil, bucket.Get(key)
	}
	return b, nil
}

func listKeys(bucket *bolt.Bucket) (keys []string, err error) {
	keys = []string{}
	err = bucket.ForEach(func(k, _ []byte) error {
		keys = append(keys, string(k))
		return nil
	})
	return
}

func (r router) ServeHTTP(w http.ResponseWriter, req *http.Request) {
	logRequest(req)
	switch req.Method {
	case "HEAD":
		getHeader(r.ctx, w, req)
	case "OPTIONS":
		w.Header().Set("Allow", "GET,PUT,DELETE,HEAD")
	case "GET":
		getBucketOrValue(r.ctx, w, req)
	case "PUT":
		putBucketOrValue(r.ctx, w, req)
	case "DELETE":
		deleteBucketOrKey(r.ctx, w, req)
	case "POST", "PATCH", "TRACE", "CONNECT":
		http.Error(w, "Method not allowed.", http.StatusMethodNotAllowed)
	default:
		http.Error(w, "Bad request.", http.StatusBadRequest)
	}
}

// returns the ETag that matched the If-None-Match
func checkIfNoneMatch(storedHeader http.Header, req *http.Request) bool {
	if nm := req.Header["If-None-Match"]; len(nm) > 0 {
		for _, m := range nm {
			if m == "*" || m == storedHeader.Get("ETag") {
				return true
			}
		}
	}
	return false
}

func getBucketOrValue(ctx context, w http.ResponseWriter, req *http.Request) {
	var (
		keys []string
		err  error
	)

	parts := [][]byte{{'/'}}
	for _, p := range bytes.Split([]byte(req.URL.EscapedPath()), []byte{'/'}) {
		if len(p) > 0 {
			parts = append(parts, p)
		}
	}

	err = ctx.db.View(func(tx *bolt.Tx) error {
		header, err := getHeaderValue(tx, req)
		if err != nil {
			return fmt.Errorf("couldn't get header: %s", err)
		}

		if header != nil {
			if checkIfNoneMatch(header, req) {
				hdr := make(http.Header)
				hdr.Set("ETag", header.Get("ETag"))
				writeHeader(hdr, w)
				w.WriteHeader(http.StatusNotModified)
				return nil
			}
		}

		p := parts
		if len(parts) > 1 {
			p = parts[:len(parts)-1]
		}
		bucket := getBoltBucket(tx, p) // get the enclosing bucket
		if bucket == nil {
			return bolt.ErrBucketNotFound
		}
		if len(parts) == 1 {
			keys, err = listKeys(bucket)
			return err
		}

		var value []byte
		bucket, value = getBoltBucketOrValue(bucket, parts[len(parts)-1])
		if bucket == nil && value == nil {
			return bolt.ErrBucketNotFound
		} else if bucket != nil {
			keys, err = listKeys(bucket)
			return err
		} else if value != nil {
			writeHeader(header, w)
			_, err := w.Write(value)
			return err
		}
		return nil
	})

	if err == bolt.ErrBucketNotFound {
		http.Error(w, "Not found", http.StatusNotFound)
		return
	} else if err != nil {
		log.Println(err)
		http.Error(w, "Out of cheese.", http.StatusInternalServerError)
		return
	}

	if keys != nil {
		w.Header().Set("Content-Type", "application/json")
		if err := json.NewEncoder(w).Encode(keys); err != nil {
			log.Println(err)
		}
		return
	}

}

func getBoltBucket(tx *bolt.Tx, parts [][]byte) *bolt.Bucket {
	b := tx.Bucket(parts[0])
	if b == nil {
		panic("nil root bucket")
	}
	for _, p := range parts[1:] {
		b = b.Bucket(p)
		if b == nil {
			return b
		}
	}
	return b
}

func getOrCreateBoltBucket(tx *bolt.Tx, parts [][]byte) (*bolt.Bucket, error) {
	b := tx.Bucket(parts[0])
	if b == nil {
		panic("nil root bucket")
	}
	var err error
	for _, p := range parts[1:] {
		b, err = b.CreateBucketIfNotExists(p)
		if err != nil {
			return nil, err
		}
	}
	return b, nil
}

func putBucketOrValue(ctx context, w http.ResponseWriter, req *http.Request) {
	parts := [][]byte{{'/'}}
	for _, p := range bytes.Split([]byte(req.URL.EscapedPath()), []byte{'/'}) {
		if len(p) > 0 {
			parts = append(parts, p)
		}
	}

	if req.ContentLength > 1<<24 {
		http.Error(w, "Request too large.", http.StatusBadRequest)
		return
	}
	key := parts[len(parts)-1]
	msg := "Out of cheese."
	status := 500
	err := ctx.db.Update(func(tx *bolt.Tx) error {
		if req.ContentLength > 0 {
			parts = parts[:len(parts)-1]
			if len(parts) == 0 {
				msg, status = "Cannot PUT a value in the root bucket.", http.StatusBadRequest
				return errors.New("request to put value in root bucket")
			}
		}
		bucket, err := getOrCreateBoltBucket(tx, parts)
		if err != nil {
			msg, status = "Error processing request.", http.StatusInternalServerError
			log.Println(err)
			return err
		}
		if req.ContentLength > 0 {
			buf := make([]byte, req.ContentLength)
			_, err := io.ReadFull(req.Body, buf)
			if err != nil && err == io.ErrUnexpectedEOF {
				msg, status = "Bad request.", http.StatusBadRequest
			}
			if err != nil {
				log.Println(err)
				return err
			}
			if err := bucket.Put(key, buf); err != nil {
				log.Println(err)
				return err
			}
			header := extractHeader(req.Header)
			eTag := etag(buf)
			header.Set("ETag", eTag)
			w.Header().Set("ETag", eTag)
			return writeHeaderValue(tx, req.URL.EscapedPath(), header)
		}
		return nil
	})
	if err != nil {
		http.Error(w, msg, status)
		return
	}
}

func writeHeaderValue(tx *bolt.Tx, path string, header http.Header) error {
	bucket := tx.Bucket(headerBucket)
	value, _ := json.Marshal(header)
	return bucket.Put([]byte(path), value)

}

func deleteBucketOrKey(ctx context, w http.ResponseWriter, req *http.Request) {
	parts := [][]byte{{'/'}}
	for _, p := range bytes.Split([]byte(req.URL.EscapedPath()), []byte{'/'}) {
		if len(p) > 0 {
			parts = append(parts, p)
		}
	}
	if len(parts) < 2 {
		http.Error(w, "Invalid path.", http.StatusBadRequest)
		return
	}
	err := ctx.db.Update(func(tx *bolt.Tx) error {
		bucket := getBoltBucket(tx, parts[:len(parts)-1])
		if bucket == nil {
			return bolt.ErrBucketNotFound
		}
		return bucket.Delete(parts[len(parts)-1])
	})
	if err == bolt.ErrBucketNotFound {
		http.Error(w, "Not found.", http.StatusNotFound)
		return
	} else if err != nil {
		http.Error(w, "Out of cheese.", http.StatusInternalServerError)
		log.Println(err)
		return
	}
}

// base64 encoded etag
func etag(b []byte) string {
	h := fnv.New64a()
	h.Write(b)
	return base64.StdEncoding.EncodeToString(h.Sum(nil))
}

func extractHeader(h http.Header) http.Header {
	result := make(http.Header)
	for _, k := range headerKeys {
		values := h[k]
		for _, v := range values {
			result.Add(k, v)
		}
	}
	return result
}

func getHeaderValue(tx *bolt.Tx, req *http.Request) (http.Header, error) {
	var header http.Header

	bucket := tx.Bucket(headerBucket)
	if bucket == nil {
		return nil, bolt.ErrBucketNotFound
	}
	h := bucket.Get([]byte(req.URL.EscapedPath()))
	if h == nil {
		return nil, nil
	}
	err := json.Unmarshal(h, &header)
	return header, err
}

func writeHeader(header http.Header, w http.ResponseWriter) {
	for key, values := range header {
		for _, value := range values {
			w.Header().Add(key, value)
		}
	}
}

func getHeader(ctx context, w http.ResponseWriter, req *http.Request) {
	var header http.Header
	err := ctx.db.View(func(tx *bolt.Tx) error {
		var err error
		header, err = getHeaderValue(tx, req)
		return err
	})
	if err == bolt.ErrBucketNotFound {
		log.Println(err)
		http.Error(w, "Header bucket not found.", http.StatusInternalServerError)
		return
	} else if err != nil {
		log.Println(err)
		http.Error(w, "Out of cheese.", http.StatusInternalServerError)
		return
	}
	if header == nil {
		http.Error(w, "Not found.", http.StatusNotFound)
		return
	}

	writeHeader(header, w)
}

func main() {
	db, err := bolt.Open(*DBName, 0600, nil)
	if err != nil {
		log.Fatal(err)
	}
	if err := createHeaderBucketIfNotExists(db); err != nil {
		log.Fatal(err)
	}
	if err := createRootBucketIfNotExists(db); err != nil {
		log.Fatal(err)
	}
	ctx := context{db}
	router := router{ctx}
	http.ListenAndServe(":8080", router)
}
<commit_msg>Add flag parsing and -port flag.<commit_after>/*
Copyright 2017 Eric Chlebek

Permission is hereby granted, free of charge, to any person obtaining a copy of
this software and associated documentation files (the "Software"), to deal in
the Software without restriction, including without limitation the rights to
use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of
the Software, and to permit persons to whom the Software is furnished to do so,
subject to the following conditions:

The above copyright notice and this permission notice shall be included in all
copies or substantial portions of the Software.

THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS
FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR
COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER
IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN
CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
*/

package main

import (
	"bytes"
	"encoding/base64"
	"encoding/json"
	"errors"
	"flag"
	"fmt"
	"hash/fnv"
	"io"
	"log"
	"net/http"

	"github.com/boltdb/bolt"
)

var (
	DBName       = flag.String("db", "bolt.db", "Bolt database to use")
	Port         = flag.Int("port", 8080, "Port to serve from")
	headerBucket = append([]byte{0}, []byte("headers")...)
	headerKeys   = []string{
		"Content-Type",
		"Content-Length",
		"ETag",
	}
)

type context struct {
	db *bolt.DB
}

type router struct {
	ctx context
}

func logRequest(req *http.Request) {
	log.Println(req.Method, req.URL.Path)
}

func createHeaderBucketIfNotExists(db *bolt.DB) error {
	return db.Update(func(tx *bolt.Tx) error {
		bucket, err := tx.CreateBucketIfNotExists(headerBucket)
		if err != nil {
			return err
		}
		return bucket.Put([]byte("/"), []byte("{}"))
	})
}

func createRootBucketIfNotExists(db *bolt.DB) error {
	return db.Update(func(tx *bolt.Tx) error {
		_, err := tx.CreateBucketIfNotExists([]byte("/"))
		return err
	})
}

func getBoltBucketOrValue(bucket *bolt.Bucket, key []byte) (*bolt.Bucket, []byte) {
	b := bucket.Bucket(key)
	if b == nil {
		return nil, bucket.Get(key)
	}
	return b, nil
}

func listKeys(bucket *bolt.Bucket) (keys []string, err error) {
	keys = []string{}
	err = bucket.ForEach(func(k, _ []byte) error {
		keys = append(keys, string(k))
		return nil
	})
	return
}

func (r router) ServeHTTP(w http.ResponseWriter, req *http.Request) {
	logRequest(req)
	switch req.Method {
	case "HEAD":
		getHeader(r.ctx, w, req)
	case "OPTIONS":
		w.Header().Set("Allow", "GET,PUT,DELETE,HEAD")
	case "GET":
		getBucketOrValue(r.ctx, w, req)
	case "PUT":
		putBucketOrValue(r.ctx, w, req)
	case "DELETE":
		deleteBucketOrKey(r.ctx, w, req)
	case "POST", "PATCH", "TRACE", "CONNECT":
		http.Error(w, "Method not allowed.", http.StatusMethodNotAllowed)
	default:
		http.Error(w, "Bad request.", http.StatusBadRequest)
	}
}

// returns the ETag that matched the If-None-Match
func checkIfNoneMatch(storedHeader http.Header, req *http.Request) bool {
	if nm := req.Header["If-None-Match"]; len(nm) > 0 {
		for _, m := range nm {
			if m == "*" || m == storedHeader.Get("ETag") {
				return true
			}
		}
	}
	return false
}

func getBucketOrValue(ctx context, w http.ResponseWriter, req *http.Request) {
	var (
		keys []string
		err  error
	)

	parts := [][]byte{{'/'}}
	for _, p := range bytes.Split([]byte(req.URL.EscapedPath()), []byte{'/'}) {
		if len(p) > 0 {
			parts = append(parts, p)
		}
	}

	err = ctx.db.View(func(tx *bolt.Tx) error {
		header, err := getHeaderValue(tx, req)
		if err != nil {
			return fmt.Errorf("couldn't get header: %s", err)
		}

		if header != nil {
			if checkIfNoneMatch(header, req) {
				hdr := make(http.Header)
				hdr.Set("ETag", header.Get("ETag"))
				writeHeader(hdr, w)
				w.WriteHeader(http.StatusNotModified)
				return nil
			}
		}

		p := parts
		if len(parts) > 1 {
			p = parts[:len(parts)-1]
		}
		bucket := getBoltBucket(tx, p) // get the enclosing bucket
		if bucket == nil {
			return bolt.ErrBucketNotFound
		}
		if len(parts) == 1 {
			keys, err = listKeys(bucket)
			return err
		}

		var value []byte
		bucket, value = getBoltBucketOrValue(bucket, parts[len(parts)-1])
		if bucket == nil && value == nil {
			return bolt.ErrBucketNotFound
		} else if bucket != nil {
			keys, err = listKeys(bucket)
			return err
		} else if value != nil {
			writeHeader(header, w)
			_, err := w.Write(value)
			return err
		}
		return nil
	})

	if err == bolt.ErrBucketNotFound {
		http.Error(w, "Not found", http.StatusNotFound)
		return
	} else if err != nil {
		log.Println(err)
		http.Error(w, "Out of cheese.", http.StatusInternalServerError)
		return
	}

	if keys != nil {
		w.Header().Set("Content-Type", "application/json")
		if err := json.NewEncoder(w).Encode(keys); err != nil {
			log.Println(err)
		}
		return
	}

}

func getBoltBucket(tx *bolt.Tx, parts [][]byte) *bolt.Bucket {
	b := tx.Bucket(parts[0])
	if b == nil {
		panic("nil root bucket")
	}
	for _, p := range parts[1:] {
		b = b.Bucket(p)
		if b == nil {
			return b
		}
	}
	return b
}

func getOrCreateBoltBucket(tx *bolt.Tx, parts [][]byte) (*bolt.Bucket, error) {
	b := tx.Bucket(parts[0])
	if b == nil {
		panic("nil root bucket")
	}
	var err error
	for _, p := range parts[1:] {
		b, err = b.CreateBucketIfNotExists(p)
		if err != nil {
			return nil, err
		}
	}
	return b, nil
}

func putBucketOrValue(ctx context, w http.ResponseWriter, req *http.Request) {
	parts := [][]byte{{'/'}}
	for _, p := range bytes.Split([]byte(req.URL.EscapedPath()), []byte{'/'}) {
		if len(p) > 0 {
			parts = append(parts, p)
		}
	}

	if req.ContentLength > 1<<24 {
		http.Error(w, "Request too large.", http.StatusBadRequest)
		return
	}
	key := parts[len(parts)-1]
	msg := "Out of cheese."
	status := 500
	err := ctx.db.Update(func(tx *bolt.Tx) error {
		if req.ContentLength > 0 {
			parts = parts[:len(parts)-1]
			if len(parts) == 0 {
				msg, status = "Cannot PUT a value in the root bucket.", http.StatusBadRequest
				return errors.New("request to put value in root bucket")
			}
		}
		bucket, err := getOrCreateBoltBucket(tx, parts)
		if err != nil {
			msg, status = "Error processing request.", http.StatusInternalServerError
			log.Println(err)
			return err
		}
		if req.ContentLength > 0 {
			buf := make([]byte, req.ContentLength)
			_, err := io.ReadFull(req.Body, buf)
			if err != nil && err == io.ErrUnexpectedEOF {
				msg, status = "Bad request.", http.StatusBadRequest
			}
			if err != nil {
				log.Println(err)
				return err
			}
			if err := bucket.Put(key, buf); err != nil {
				log.Println(err)
				return err
			}
			header := extractHeader(req.Header)
			eTag := etag(buf)
			header.Set("ETag", eTag)
			w.Header().Set("ETag", eTag)
			return writeHeaderValue(tx, req.URL.EscapedPath(), header)
		}
		return nil
	})
	if err != nil {
		http.Error(w, msg, status)
		return
	}
}

func writeHeaderValue(tx *bolt.Tx, path string, header http.Header) error {
	bucket := tx.Bucket(headerBucket)
	value, _ := json.Marshal(header)
	return bucket.Put([]byte(path), value)

}

func deleteBucketOrKey(ctx context, w http.ResponseWriter, req *http.Request) {
	parts := [][]byte{{'/'}}
	for _, p := range bytes.Split([]byte(req.URL.EscapedPath()), []byte{'/'}) {
		if len(p) > 0 {
			parts = append(parts, p)
		}
	}
	if len(parts) < 2 {
		http.Error(w, "Invalid path.", http.StatusBadRequest)
		return
	}
	err := ctx.db.Update(func(tx *bolt.Tx) error {
		bucket := getBoltBucket(tx, parts[:len(parts)-1])
		if bucket == nil {
			return bolt.ErrBucketNotFound
		}
		return bucket.Delete(parts[len(parts)-1])
	})
	if err == bolt.ErrBucketNotFound {
		http.Error(w, "Not found.", http.StatusNotFound)
		return
	} else if err != nil {
		http.Error(w, "Out of cheese.", http.StatusInternalServerError)
		log.Println(err)
		return
	}
}

// base64 encoded etag
func etag(b []byte) string {
	h := fnv.New64a()
	h.Write(b)
	return base64.StdEncoding.EncodeToString(h.Sum(nil))
}

func extractHeader(h http.Header) http.Header {
	result := make(http.Header)
	for _, k := range headerKeys {
		values := h[k]
		for _, v := range values {
			result.Add(k, v)
		}
	}
	return result
}

func getHeaderValue(tx *bolt.Tx, req *http.Request) (http.Header, error) {
	var header http.Header

	bucket := tx.Bucket(headerBucket)
	if bucket == nil {
		return nil, bolt.ErrBucketNotFound
	}
	h := bucket.Get([]byte(req.URL.EscapedPath()))
	if h == nil {
		return nil, nil
	}
	err := json.Unmarshal(h, &header)
	return header, err
}

func writeHeader(header http.Header, w http.ResponseWriter) {
	for key, values := range header {
		for _, value := range values {
			w.Header().Add(key, value)
		}
	}
}

func getHeader(ctx context, w http.ResponseWriter, req *http.Request) {
	var header http.Header
	err := ctx.db.View(func(tx *bolt.Tx) error {
		var err error
		header, err = getHeaderValue(tx, req)
		return err
	})
	if err == bolt.ErrBucketNotFound {
		log.Println(err)
		http.Error(w, "Header bucket not found.", http.StatusInternalServerError)
		return
	} else if err != nil {
		log.Println(err)
		http.Error(w, "Out of cheese.", http.StatusInternalServerError)
		return
	}
	if header == nil {
		http.Error(w, "Not found.", http.StatusNotFound)
		return
	}

	writeHeader(header, w)
}

func main() {
	flag.Parse()
	db, err := bolt.Open(*DBName, 0600, nil)
	if err != nil {
		log.Fatal(err)
	}
	if err := createHeaderBucketIfNotExists(db); err != nil {
		log.Fatal(err)
	}
	if err := createRootBucketIfNotExists(db); err != nil {
		log.Fatal(err)
	}
	ctx := context{db}
	router := router{ctx}
	http.ListenAndServe(fmt.Sprintf(":%d", *Port), router)
}
<|endoftext|>


<commit_before>package devd

import (
	"crypto/tls"
	"fmt"
	"html/template"
	"net"
	"net/http"
	"os"
	"os/signal"
	"regexp"
	"strings"
	"syscall"
	"time"

	"golang.org/x/net/context"

	"github.com/GeertJohan/go.rice"
	"github.com/goji/httpauth"

	"github.com/cortesi/devd/httpctx"
	"github.com/cortesi/devd/inject"
	"github.com/cortesi/devd/livereload"
	"github.com/cortesi/devd/ricetemp"
	"github.com/cortesi/devd/slowdown"
	"github.com/cortesi/devd/timer"
	"github.com/cortesi/termlog"
)

const (
	// Version is the current version of devd
	Version  = "0.5.1"
	portLow  = 8000
	portHigh = 10000
)

func pickPort(addr string, low int, high int, tls bool) (net.Listener, error) {
	firstTry := 80
	if tls {
		firstTry = 443
	}
	hl, err := net.Listen("tcp", fmt.Sprintf("%v:%d", addr, firstTry))
	if err == nil {
		return hl, nil
	}
	for i := low; i < high; i++ {
		hl, err := net.Listen("tcp", fmt.Sprintf("%v:%d", addr, i))
		if err == nil {
			return hl, nil
		}
	}
	return nil, fmt.Errorf("Could not find open port.")
}

func getTLSConfig(path string) (t *tls.Config, err error) {
	config := &tls.Config{}
	if config.NextProtos == nil {
		config.NextProtos = []string{"http/1.1"}
	}
	config.Certificates = make([]tls.Certificate, 1)
	config.Certificates[0], err = tls.LoadX509KeyPair(path, path)
	if err != nil {
		return nil, err
	}
	return config, nil
}

// This filthy hack works in conjunction with hostPortStrip to restore the
// original request host after mux match.
func revertOriginalHost(r *http.Request) {
	original := r.Header.Get("_devd_original_host")
	if original != "" {
		r.Host = original
		r.Header.Del("_devd_original_host")
	}
}

// We can remove the mangling once this is fixed:
// 		https://github.com/golang/go/issues/10463
func hostPortStrip(next http.Handler) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		host, _, err := net.SplitHostPort(r.Host)
		if err == nil {
			original := r.Host
			r.Host = host
			r.Header.Set("_devd_original_host", original)
		}
		next.ServeHTTP(w, r)
	})
}

func matchStringAny(regexps []*regexp.Regexp, s string) bool {
	for _, r := range regexps {
		if r.MatchString(s) {
			return true
		}
	}
	return false
}

func formatURL(tls bool, httpIP string, port int) string {
	proto := "http"
	if tls {
		proto = "https"
	}
	host := httpIP
	if httpIP == "0.0.0.0" || httpIP == "127.0.0.1" {
		host = "devd.io"
	}
	if port == 443 && tls {
		return fmt.Sprintf("https://%s", host)
	}
	if port == 80 && !tls {
		return fmt.Sprintf("http://%s", host)
	}
	return fmt.Sprintf("%s://%s:%d", proto, host, port)
}

// Credentials is a simple username/password pair
type Credentials struct {
	username string
	password string
}

// CredentialsFromSpec creates a set of credentials from a spec
func CredentialsFromSpec(spec string) (*Credentials, error) {
	parts := strings.SplitN(spec, ":", 2)
	if len(parts) != 2 || parts[0] == "" || parts[1] == "" {
		return nil, fmt.Errorf("Invalid credential spec: %s", spec)
	}
	return &Credentials{parts[0], parts[1]}, nil
}

// Devd represents the devd server options
type Devd struct {
	Routes RouteCollection

	// Shaping
	Latency  int
	DownKbps uint
	UpKbps   uint

	// Add headers
	AddHeaders *http.Header

	// Livereload and watch static routes
	LivereloadRoutes bool
	// Livereload, but don't watch static routes
	Livereload bool
	WatchPaths []string
	Excludes   []string

	// Logging
	IgnoreLogs []*regexp.Regexp

	// Password protection
	Credentials *Credentials

	lrserver *livereload.Server
}

// WrapHandler wraps an httpctx.Handler in the paraphernalia needed by devd for
// logging, latency, and so forth.
func (dd *Devd) WrapHandler(log termlog.TermLog, next httpctx.Handler) http.Handler {
	h := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		revertOriginalHost(r)
		timr := timer.Timer{}
		sublog := log.Group()
		defer func() {
			timing := termlog.DefaultPalette.Timestamp.SprintFunc()("timing: ")
			sublog.SayAs("timer", timing+timr.String())
			sublog.Done()
		}()
		if matchStringAny(dd.IgnoreLogs, fmt.Sprintf("%s%s", r.URL.Host, r.RequestURI)) {
			sublog.Quiet()
		}
		timr.RequestHeaders()
		time.Sleep(time.Millisecond * time.Duration(dd.Latency))

		dpath := r.URL.String()
		if !strings.HasPrefix(dpath, "/") {
			dpath = "/" + dpath
		}
		sublog.Say("%s %s", r.Method, dpath)
		LogHeader(sublog, r.Header)
		ctx := timr.NewContext(context.Background())
		ctx = termlog.NewContext(ctx, sublog)
		if dd.AddHeaders != nil {
			for h, vals := range *dd.AddHeaders {
				for _, v := range vals {
					w.Header().Set(h, v)
				}
			}
		}
		next.ServeHTTPContext(
			ctx,
			&ResponseLogWriter{Log: sublog, Resp: w, Timer: &timr},
			r,
		)
	})
	return h
}

// HasLivereload tells us if livereload is enabled
func (dd *Devd) HasLivereload() bool {
	if dd.Livereload || dd.LivereloadRoutes || len(dd.WatchPaths) > 0 {
		return true
	}
	return false
}

// AddRoutes adds route specifications to the server
func (dd *Devd) AddRoutes(specs []string) error {
	dd.Routes = make(RouteCollection)
	for _, s := range specs {
		err := dd.Routes.Add(s)
		if err != nil {
			return fmt.Errorf("Invalid route specification: %s", err)
		}
	}
	return nil
}

// AddIgnores adds log ignore patterns to the server
func (dd *Devd) AddIgnores(specs []string) error {
	dd.IgnoreLogs = make([]*regexp.Regexp, 0, 0)
	for _, expr := range specs {
		v, err := regexp.Compile(expr)
		if err != nil {
			return fmt.Errorf("%s", err)
		}
		dd.IgnoreLogs = append(dd.IgnoreLogs, v)
	}
	return nil
}

// HandleNotFound handles pages not found. In particular, this handler is used
// when we have no matching route for a request. This also means it's not
// useful to inject the livereload paraphernalia here.
func HandleNotFound(templates *template.Template) httpctx.Handler {
	return httpctx.HandlerFunc(func(ctx context.Context, w http.ResponseWriter, r *http.Request) {
		w.WriteHeader(http.StatusNotFound)
		err := templates.Lookup("404.html").Execute(w, nil)
		if err != nil {
			logger := termlog.FromContext(ctx)
			logger.Shout("Could not execute template: %s", err)
		}
	})
}

// Router constructs the main Devd router that serves all requests
func (dd *Devd) Router(logger termlog.TermLog, templates *template.Template) (http.Handler, error) {
	mux := http.NewServeMux()
	hasGlobal := false

	ci := inject.CopyInject{}
	if dd.HasLivereload() {
		ci = livereload.Injector
	}

	for match, route := range dd.Routes {
		if match == "/" {
			hasGlobal = true
		}
		handler := dd.WrapHandler(
			logger,
			route.Endpoint.Handler(templates, ci),
		)
		handler = http.StripPrefix(route.Path, handler)
		mux.Handle(match, handler)
	}
	if dd.HasLivereload() {
		lr := livereload.NewServer("livereload", logger)
		mux.Handle(livereload.EndpointPath, lr)
		mux.Handle(livereload.ScriptPath, http.HandlerFunc(lr.ServeScript))
		if dd.LivereloadRoutes {
			err := WatchRoutes(dd.Routes, lr, dd.Excludes, logger)
			if err != nil {
				return nil, fmt.Errorf("Could not watch routes for livereload: %s", err)
			}
		}
		if len(dd.WatchPaths) > 0 {
			err := WatchPaths(dd.WatchPaths, dd.Excludes, lr, logger)
			if err != nil {
				return nil, fmt.Errorf("Could not watch path for livereload: %s", err)
			}
		}
		dd.lrserver = lr
	}
	if !hasGlobal {
		mux.Handle(
			"/",
			dd.WrapHandler(logger, HandleNotFound(templates)),
		)
	}
	var h = http.Handler(mux)
	if dd.Credentials != nil {
		h = httpauth.SimpleBasicAuth(
			dd.Credentials.username, dd.Credentials.password,
		)(h)
	}
	return hostPortStrip(h), nil
}

// Serve starts the devd server. The callback is called with the serving URL
// just before service starts.
func (dd *Devd) Serve(address string, port int, certFile string, logger termlog.TermLog, callback func(string)) error {
	templates, err := ricetemp.MakeTemplates(rice.MustFindBox("templates"))
	if err != nil {
		return fmt.Errorf("Error loading templates: %s", err)
	}
	mux, err := dd.Router(logger, templates)
	if err != nil {
		return err
	}
	var tlsConfig *tls.Config
	var tlsEnabled bool
	if certFile != "" {
		tlsConfig, err = getTLSConfig(certFile)
		if err != nil {
			return fmt.Errorf("Could not load certs: %s", err)
		}
		tlsEnabled = true
	}

	var hl net.Listener
	if port > 0 {
		hl, err = net.Listen("tcp", fmt.Sprintf("%v:%d", address, port))
	} else {
		hl, err = pickPort(address, portLow, portHigh, tlsEnabled)
	}
	if err != nil {
		return err
	}

	if tlsConfig != nil {
		hl = tls.NewListener(hl, tlsConfig)
	}

	hl = slowdown.NewSlowListener(hl, dd.UpKbps*1024, dd.DownKbps*1024)
	url := formatURL(tlsEnabled, address, hl.Addr().(*net.TCPAddr).Port)
	logger.Say("Listening on %s (%s)", url, hl.Addr().String())
	server := &http.Server{Addr: hl.Addr().String(), Handler: mux}
	callback(url)

	if dd.HasLivereload() {
		c := make(chan os.Signal, 1)
		signal.Notify(c, syscall.SIGHUP)
		go func() {
			for {
				<-c
				logger.Say("Received signal - reloading")
				dd.lrserver.Reload([]string{"*"})
			}
		}()
	}

	err = server.Serve(hl)
	logger.Shout("Server stopped: %v", err)
	return nil
}
<commit_msg>Run livereload endpoints on all domains<commit_after>package devd

import (
	"crypto/tls"
	"fmt"
	"html/template"
	"net"
	"net/http"
	"os"
	"os/signal"
	"regexp"
	"strings"
	"syscall"
	"time"

	"golang.org/x/net/context"

	"github.com/GeertJohan/go.rice"
	"github.com/goji/httpauth"

	"github.com/cortesi/devd/httpctx"
	"github.com/cortesi/devd/inject"
	"github.com/cortesi/devd/livereload"
	"github.com/cortesi/devd/ricetemp"
	"github.com/cortesi/devd/slowdown"
	"github.com/cortesi/devd/timer"
	"github.com/cortesi/termlog"
)

const (
	// Version is the current version of devd
	Version  = "0.5.1"
	portLow  = 8000
	portHigh = 10000
)

func pickPort(addr string, low int, high int, tls bool) (net.Listener, error) {
	firstTry := 80
	if tls {
		firstTry = 443
	}
	hl, err := net.Listen("tcp", fmt.Sprintf("%v:%d", addr, firstTry))
	if err == nil {
		return hl, nil
	}
	for i := low; i < high; i++ {
		hl, err := net.Listen("tcp", fmt.Sprintf("%v:%d", addr, i))
		if err == nil {
			return hl, nil
		}
	}
	return nil, fmt.Errorf("Could not find open port.")
}

func getTLSConfig(path string) (t *tls.Config, err error) {
	config := &tls.Config{}
	if config.NextProtos == nil {
		config.NextProtos = []string{"http/1.1"}
	}
	config.Certificates = make([]tls.Certificate, 1)
	config.Certificates[0], err = tls.LoadX509KeyPair(path, path)
	if err != nil {
		return nil, err
	}
	return config, nil
}

// This filthy hack works in conjunction with hostPortStrip to restore the
// original request host after mux match.
func revertOriginalHost(r *http.Request) {
	original := r.Header.Get("_devd_original_host")
	if original != "" {
		r.Host = original
		r.Header.Del("_devd_original_host")
	}
}

// We can remove the mangling once this is fixed:
// 		https://github.com/golang/go/issues/10463
func hostPortStrip(next http.Handler) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		host, _, err := net.SplitHostPort(r.Host)
		if err == nil {
			original := r.Host
			r.Host = host
			r.Header.Set("_devd_original_host", original)
		}
		next.ServeHTTP(w, r)
	})
}

func matchStringAny(regexps []*regexp.Regexp, s string) bool {
	for _, r := range regexps {
		if r.MatchString(s) {
			return true
		}
	}
	return false
}

func formatURL(tls bool, httpIP string, port int) string {
	proto := "http"
	if tls {
		proto = "https"
	}
	host := httpIP
	if httpIP == "0.0.0.0" || httpIP == "127.0.0.1" {
		host = "devd.io"
	}
	if port == 443 && tls {
		return fmt.Sprintf("https://%s", host)
	}
	if port == 80 && !tls {
		return fmt.Sprintf("http://%s", host)
	}
	return fmt.Sprintf("%s://%s:%d", proto, host, port)
}

// Credentials is a simple username/password pair
type Credentials struct {
	username string
	password string
}

// CredentialsFromSpec creates a set of credentials from a spec
func CredentialsFromSpec(spec string) (*Credentials, error) {
	parts := strings.SplitN(spec, ":", 2)
	if len(parts) != 2 || parts[0] == "" || parts[1] == "" {
		return nil, fmt.Errorf("Invalid credential spec: %s", spec)
	}
	return &Credentials{parts[0], parts[1]}, nil
}

// Devd represents the devd server options
type Devd struct {
	Routes RouteCollection

	// Shaping
	Latency  int
	DownKbps uint
	UpKbps   uint

	// Add headers
	AddHeaders *http.Header

	// Livereload and watch static routes
	LivereloadRoutes bool
	// Livereload, but don't watch static routes
	Livereload bool
	WatchPaths []string
	Excludes   []string

	// Logging
	IgnoreLogs []*regexp.Regexp

	// Password protection
	Credentials *Credentials

	lrserver *livereload.Server
}

// WrapHandler wraps an httpctx.Handler in the paraphernalia needed by devd for
// logging, latency, and so forth.
func (dd *Devd) WrapHandler(log termlog.TermLog, next httpctx.Handler) http.Handler {
	h := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		revertOriginalHost(r)
		timr := timer.Timer{}
		sublog := log.Group()
		defer func() {
			timing := termlog.DefaultPalette.Timestamp.SprintFunc()("timing: ")
			sublog.SayAs("timer", timing+timr.String())
			sublog.Done()
		}()
		if matchStringAny(dd.IgnoreLogs, fmt.Sprintf("%s%s", r.URL.Host, r.RequestURI)) {
			sublog.Quiet()
		}
		timr.RequestHeaders()
		time.Sleep(time.Millisecond * time.Duration(dd.Latency))

		dpath := r.URL.String()
		if !strings.HasPrefix(dpath, "/") {
			dpath = "/" + dpath
		}
		sublog.Say("%s %s", r.Method, dpath)
		LogHeader(sublog, r.Header)
		ctx := timr.NewContext(context.Background())
		ctx = termlog.NewContext(ctx, sublog)
		if dd.AddHeaders != nil {
			for h, vals := range *dd.AddHeaders {
				for _, v := range vals {
					w.Header().Set(h, v)
				}
			}
		}
		next.ServeHTTPContext(
			ctx,
			&ResponseLogWriter{Log: sublog, Resp: w, Timer: &timr},
			r,
		)
	})
	return h
}

// HasLivereload tells us if livereload is enabled
func (dd *Devd) HasLivereload() bool {
	if dd.Livereload || dd.LivereloadRoutes || len(dd.WatchPaths) > 0 {
		return true
	}
	return false
}

// AddRoutes adds route specifications to the server
func (dd *Devd) AddRoutes(specs []string) error {
	dd.Routes = make(RouteCollection)
	for _, s := range specs {
		err := dd.Routes.Add(s)
		if err != nil {
			return fmt.Errorf("Invalid route specification: %s", err)
		}
	}
	return nil
}

// AddIgnores adds log ignore patterns to the server
func (dd *Devd) AddIgnores(specs []string) error {
	dd.IgnoreLogs = make([]*regexp.Regexp, 0, 0)
	for _, expr := range specs {
		v, err := regexp.Compile(expr)
		if err != nil {
			return fmt.Errorf("%s", err)
		}
		dd.IgnoreLogs = append(dd.IgnoreLogs, v)
	}
	return nil
}

// HandleNotFound handles pages not found. In particular, this handler is used
// when we have no matching route for a request. This also means it's not
// useful to inject the livereload paraphernalia here.
func HandleNotFound(templates *template.Template) httpctx.Handler {
	return httpctx.HandlerFunc(func(ctx context.Context, w http.ResponseWriter, r *http.Request) {
		w.WriteHeader(http.StatusNotFound)
		err := templates.Lookup("404.html").Execute(w, nil)
		if err != nil {
			logger := termlog.FromContext(ctx)
			logger.Shout("Could not execute template: %s", err)
		}
	})
}

// Router constructs the main Devd router that serves all requests
func (dd *Devd) Router(logger termlog.TermLog, templates *template.Template) (http.Handler, error) {
	mux := http.NewServeMux()
	hasGlobal := false

	ci := inject.CopyInject{}
	if dd.HasLivereload() {
		ci = livereload.Injector
	}

	for match, route := range dd.Routes {
		if match == "/" {
			hasGlobal = true
		}
		handler := dd.WrapHandler(
			logger,
			route.Endpoint.Handler(templates, ci),
		)
		handler = http.StripPrefix(route.Path, handler)
		mux.Handle(match, handler)
	}
	if dd.HasLivereload() {
		lr := livereload.NewServer("livereload", logger)
		mux.Handle(livereload.EndpointPath, lr)
		mux.Handle(livereload.ScriptPath, http.HandlerFunc(lr.ServeScript))
		seen := make(map[string]bool)
		for _, route := range dd.Routes {
			if _, ok := seen[route.Host]; route.Host != "" && ok == false {
				mux.Handle(route.Host+livereload.EndpointPath, lr)
				mux.Handle(
					route.Host+livereload.ScriptPath,
					http.HandlerFunc(lr.ServeScript),
				)
				seen[route.Host] = true
			}
		}
		if dd.LivereloadRoutes {
			err := WatchRoutes(dd.Routes, lr, dd.Excludes, logger)
			if err != nil {
				return nil, fmt.Errorf("Could not watch routes for livereload: %s", err)
			}
		}
		if len(dd.WatchPaths) > 0 {
			err := WatchPaths(dd.WatchPaths, dd.Excludes, lr, logger)
			if err != nil {
				return nil, fmt.Errorf("Could not watch path for livereload: %s", err)
			}
		}
		dd.lrserver = lr
	}
	if !hasGlobal {
		mux.Handle(
			"/",
			dd.WrapHandler(logger, HandleNotFound(templates)),
		)
	}
	var h = http.Handler(mux)
	if dd.Credentials != nil {
		h = httpauth.SimpleBasicAuth(
			dd.Credentials.username, dd.Credentials.password,
		)(h)
	}
	return hostPortStrip(h), nil
}

// Serve starts the devd server. The callback is called with the serving URL
// just before service starts.
func (dd *Devd) Serve(address string, port int, certFile string, logger termlog.TermLog, callback func(string)) error {
	templates, err := ricetemp.MakeTemplates(rice.MustFindBox("templates"))
	if err != nil {
		return fmt.Errorf("Error loading templates: %s", err)
	}
	mux, err := dd.Router(logger, templates)
	if err != nil {
		return err
	}
	var tlsConfig *tls.Config
	var tlsEnabled bool
	if certFile != "" {
		tlsConfig, err = getTLSConfig(certFile)
		if err != nil {
			return fmt.Errorf("Could not load certs: %s", err)
		}
		tlsEnabled = true
	}

	var hl net.Listener
	if port > 0 {
		hl, err = net.Listen("tcp", fmt.Sprintf("%v:%d", address, port))
	} else {
		hl, err = pickPort(address, portLow, portHigh, tlsEnabled)
	}
	if err != nil {
		return err
	}

	if tlsConfig != nil {
		hl = tls.NewListener(hl, tlsConfig)
	}

	hl = slowdown.NewSlowListener(hl, dd.UpKbps*1024, dd.DownKbps*1024)
	url := formatURL(tlsEnabled, address, hl.Addr().(*net.TCPAddr).Port)
	logger.Say("Listening on %s (%s)", url, hl.Addr().String())
	server := &http.Server{Addr: hl.Addr().String(), Handler: mux}
	callback(url)

	if dd.HasLivereload() {
		c := make(chan os.Signal, 1)
		signal.Notify(c, syscall.SIGHUP)
		go func() {
			for {
				<-c
				logger.Say("Received signal - reloading")
				dd.lrserver.Reload([]string{"*"})
			}
		}()
	}

	err = server.Serve(hl)
	logger.Shout("Server stopped: %v", err)
	return nil
}
<|endoftext|>


<commit_before>package main

import (
	"bytes"
	"context"
	"crypto/subtle"
	"fmt"
	"log"
	"net"
	"net/http"
	"strconv"
	"strings"
	"sync"
	"time"

	nanoid "github.com/matoous/go-nanoid"
	"golang.org/x/net/netutil"
)

const healthPath = "/health"

var (
	mimes = map[imageType]string{
		imageTypeJPEG: "image/jpeg",
		imageTypePNG:  "image/png",
		imageTypeWEBP: "image/webp",
	}

	authHeaderMust []byte

	imgproxyIsRunningMsg = []byte("imgproxy is running")

	errInvalidMethod = newError(422, "Invalid request method", "Method doesn't allowed")
	errInvalidSecret = newError(403, "Invalid secret", "Forbidden")
)

var responseBufPool = sync.Pool{
	New: func() interface{} {
		return new(bytes.Buffer)
	},
}

type httpHandler struct {
	sem chan struct{}
}

func newHTTPHandler() *httpHandler {
	return &httpHandler{make(chan struct{}, conf.Concurrency)}
}

func startServer() *http.Server {
	l, err := net.Listen("tcp", conf.Bind)
	if err != nil {
		log.Fatal(err)
	}
	s := &http.Server{
		Handler:        newHTTPHandler(),
		ReadTimeout:    time.Duration(conf.ReadTimeout) * time.Second,
		MaxHeaderBytes: 1 << 20,
	}

	go func() {
		log.Printf("Starting server at %s\n", conf.Bind)
		if err := s.Serve(netutil.LimitListener(l, conf.MaxClients)); err != nil && err != http.ErrServerClosed {
			log.Fatalln(err)
		}
	}()

	return s
}

func shutdownServer(s *http.Server) {
	log.Println("Shutting down the server...")

	ctx, close := context.WithTimeout(context.Background(), 5*time.Second)
	defer close()

	s.Shutdown(ctx)
}

func logResponse(status int, msg string) {
	var color int

	if status >= 500 {
		color = 31
	} else if status >= 400 {
		color = 33
	} else {
		color = 32
	}

	log.Printf("|\033[7;%dm %d \033[0m| %s\n", color, status, msg)
}

func writeCORS(rw http.ResponseWriter) {
	if len(conf.AllowOrigin) > 0 {
		rw.Header().Set("Access-Control-Allow-Origin", conf.AllowOrigin)
		rw.Header().Set("Access-Control-Allow-Methods", "GET, OPTIONs")
	}
}

func respondWithImage(ctx context.Context, reqID string, r *http.Request, rw http.ResponseWriter, data []byte) {
	po := getProcessingOptions(ctx)

	rw.Header().Set("Expires", time.Now().Add(time.Second*time.Duration(conf.TTL)).Format(http.TimeFormat))
	rw.Header().Set("Cache-Control", fmt.Sprintf("max-age=%d, public", conf.TTL))
	rw.Header().Set("Content-Type", mimes[po.Format])

	dataToRespond := data

	if conf.GZipCompression > 0 && strings.Contains(r.Header.Get("Accept-Encoding"), "gzip") {
		rw.Header().Set("Content-Encoding", "gzip")

		buf := responseBufPool.Get().(*bytes.Buffer)
		defer responseBufPool.Put(buf)

		gzipData(data, buf)
		dataToRespond = buf.Bytes()
	}

	rw.Header().Set("Content-Length", strconv.Itoa(len(dataToRespond)))
	rw.WriteHeader(200)
	rw.Write(dataToRespond)

	logResponse(200, fmt.Sprintf("[%s] Processed in %s: %s; %+v", reqID, getTimerSince(ctx), getImageURL(ctx), po))
}

func respondWithError(reqID string, rw http.ResponseWriter, err imgproxyError) {
	logResponse(err.StatusCode, fmt.Sprintf("[%s] %s", reqID, err.Message))

	rw.WriteHeader(err.StatusCode)
	rw.Write([]byte(err.PublicMessage))
}

func respondWithOptions(reqID string, rw http.ResponseWriter) {
	logResponse(200, fmt.Sprintf("[%s] Respond with options", reqID))
	rw.WriteHeader(200)
}

func prepareAuthHeaderMust() []byte {
	if len(authHeaderMust) == 0 {
		authHeaderMust = []byte(fmt.Sprintf("Bearer %s", conf.Secret))
	}

	return authHeaderMust
}

func checkSecret(r *http.Request) bool {
	if len(conf.Secret) == 0 {
		return true
	}

	return subtle.ConstantTimeCompare(
		[]byte(r.Header.Get("Authorization")),
		prepareAuthHeaderMust(),
	) == 1
}

func (h *httpHandler) lock() {
	h.sem <- struct{}{}
}

func (h *httpHandler) unlock() {
	<-h.sem
}

func (h *httpHandler) ServeHTTP(rw http.ResponseWriter, r *http.Request) {
	reqID, _ := nanoid.Nanoid()

	defer func() {
		if r := recover(); r != nil {
			if err, ok := r.(imgproxyError); ok {
				respondWithError(reqID, rw, err)
			} else {
				respondWithError(reqID, rw, newUnexpectedError(r.(error), 4))
			}
		}
	}()

	log.Printf("[%s] %s: %s\n", reqID, r.Method, r.URL.RequestURI())

	writeCORS(rw)

	if r.Method == http.MethodOptions {
		respondWithOptions(reqID, rw)
		return
	}

	if r.Method != http.MethodGet {
		panic(errInvalidMethod)
	}

	if !checkSecret(r) {
		panic(errInvalidSecret)
	}

	ctx := context.Background()

	if newRelicEnabled {
		var newRelicCancel context.CancelFunc
		ctx, newRelicCancel = startNewRelicTransaction(ctx, rw, r)
		defer newRelicCancel()
	}

	h.lock()
	defer h.unlock()

	if r.URL.RequestURI() == healthPath {
		rw.WriteHeader(200)
		rw.Write(imgproxyIsRunningMsg)
		return
	}

	ctx, timeoutCancel := startTimer(ctx, time.Duration(conf.WriteTimeout)*time.Second)
	defer timeoutCancel()

	ctx, err := parsePath(ctx, r)
	if err != nil {
		panic(newError(404, err.Error(), "Invalid image url"))
	}

	ctx, downloadcancel, err := downloadImage(ctx)
	defer downloadcancel()
	if err != nil {
		if newRelicEnabled {
			sendErrorToNewRelic(ctx, err)
		}
		panic(newError(404, err.Error(), "Image is unreachable"))
	}

	checkTimeout(ctx)

	if conf.ETagEnabled {
		eTag := calcETag(ctx)
		rw.Header().Set("ETag", eTag)

		if eTag == r.Header.Get("If-None-Match") {
			panic(errNotModified)
		}
	}

	checkTimeout(ctx)

	imageData, err := processImage(ctx)
	if err != nil {
		if newRelicEnabled {
			sendErrorToNewRelic(ctx, err)
		}
		panic(newError(500, err.Error(), "Error occurred while processing image"))
	}

	checkTimeout(ctx)

	respondWithImage(ctx, reqID, r, rw, imageData)
}
<commit_msg>Reset response buffer<commit_after>package main

import (
	"bytes"
	"context"
	"crypto/subtle"
	"fmt"
	"log"
	"net"
	"net/http"
	"strconv"
	"strings"
	"sync"
	"time"

	nanoid "github.com/matoous/go-nanoid"
	"golang.org/x/net/netutil"
)

const healthPath = "/health"

var (
	mimes = map[imageType]string{
		imageTypeJPEG: "image/jpeg",
		imageTypePNG:  "image/png",
		imageTypeWEBP: "image/webp",
	}

	authHeaderMust []byte

	imgproxyIsRunningMsg = []byte("imgproxy is running")

	errInvalidMethod = newError(422, "Invalid request method", "Method doesn't allowed")
	errInvalidSecret = newError(403, "Invalid secret", "Forbidden")
)

var responseBufPool = sync.Pool{
	New: func() interface{} {
		return new(bytes.Buffer)
	},
}

type httpHandler struct {
	sem chan struct{}
}

func newHTTPHandler() *httpHandler {
	return &httpHandler{make(chan struct{}, conf.Concurrency)}
}

func startServer() *http.Server {
	l, err := net.Listen("tcp", conf.Bind)
	if err != nil {
		log.Fatal(err)
	}
	s := &http.Server{
		Handler:        newHTTPHandler(),
		ReadTimeout:    time.Duration(conf.ReadTimeout) * time.Second,
		MaxHeaderBytes: 1 << 20,
	}

	go func() {
		log.Printf("Starting server at %s\n", conf.Bind)
		if err := s.Serve(netutil.LimitListener(l, conf.MaxClients)); err != nil && err != http.ErrServerClosed {
			log.Fatalln(err)
		}
	}()

	return s
}

func shutdownServer(s *http.Server) {
	log.Println("Shutting down the server...")

	ctx, close := context.WithTimeout(context.Background(), 5*time.Second)
	defer close()

	s.Shutdown(ctx)
}

func logResponse(status int, msg string) {
	var color int

	if status >= 500 {
		color = 31
	} else if status >= 400 {
		color = 33
	} else {
		color = 32
	}

	log.Printf("|\033[7;%dm %d \033[0m| %s\n", color, status, msg)
}

func writeCORS(rw http.ResponseWriter) {
	if len(conf.AllowOrigin) > 0 {
		rw.Header().Set("Access-Control-Allow-Origin", conf.AllowOrigin)
		rw.Header().Set("Access-Control-Allow-Methods", "GET, OPTIONs")
	}
}

func respondWithImage(ctx context.Context, reqID string, r *http.Request, rw http.ResponseWriter, data []byte) {
	po := getProcessingOptions(ctx)

	rw.Header().Set("Expires", time.Now().Add(time.Second*time.Duration(conf.TTL)).Format(http.TimeFormat))
	rw.Header().Set("Cache-Control", fmt.Sprintf("max-age=%d, public", conf.TTL))
	rw.Header().Set("Content-Type", mimes[po.Format])

	dataToRespond := data

	if conf.GZipCompression > 0 && strings.Contains(r.Header.Get("Accept-Encoding"), "gzip") {
		rw.Header().Set("Content-Encoding", "gzip")

		buf := responseBufPool.Get().(*bytes.Buffer)
		buf.Reset()
		defer responseBufPool.Put(buf)

		gzipData(data, buf)
		dataToRespond = buf.Bytes()
	}

	rw.Header().Set("Content-Length", strconv.Itoa(len(dataToRespond)))
	rw.WriteHeader(200)
	rw.Write(dataToRespond)

	logResponse(200, fmt.Sprintf("[%s] Processed in %s: %s; %+v", reqID, getTimerSince(ctx), getImageURL(ctx), po))
}

func respondWithError(reqID string, rw http.ResponseWriter, err imgproxyError) {
	logResponse(err.StatusCode, fmt.Sprintf("[%s] %s", reqID, err.Message))

	rw.WriteHeader(err.StatusCode)
	rw.Write([]byte(err.PublicMessage))
}

func respondWithOptions(reqID string, rw http.ResponseWriter) {
	logResponse(200, fmt.Sprintf("[%s] Respond with options", reqID))
	rw.WriteHeader(200)
}

func prepareAuthHeaderMust() []byte {
	if len(authHeaderMust) == 0 {
		authHeaderMust = []byte(fmt.Sprintf("Bearer %s", conf.Secret))
	}

	return authHeaderMust
}

func checkSecret(r *http.Request) bool {
	if len(conf.Secret) == 0 {
		return true
	}

	return subtle.ConstantTimeCompare(
		[]byte(r.Header.Get("Authorization")),
		prepareAuthHeaderMust(),
	) == 1
}

func (h *httpHandler) lock() {
	h.sem <- struct{}{}
}

func (h *httpHandler) unlock() {
	<-h.sem
}

func (h *httpHandler) ServeHTTP(rw http.ResponseWriter, r *http.Request) {
	reqID, _ := nanoid.Nanoid()

	defer func() {
		if r := recover(); r != nil {
			if err, ok := r.(imgproxyError); ok {
				respondWithError(reqID, rw, err)
			} else {
				respondWithError(reqID, rw, newUnexpectedError(r.(error), 4))
			}
		}
	}()

	log.Printf("[%s] %s: %s\n", reqID, r.Method, r.URL.RequestURI())

	writeCORS(rw)

	if r.Method == http.MethodOptions {
		respondWithOptions(reqID, rw)
		return
	}

	if r.Method != http.MethodGet {
		panic(errInvalidMethod)
	}

	if !checkSecret(r) {
		panic(errInvalidSecret)
	}

	ctx := context.Background()

	if newRelicEnabled {
		var newRelicCancel context.CancelFunc
		ctx, newRelicCancel = startNewRelicTransaction(ctx, rw, r)
		defer newRelicCancel()
	}

	h.lock()
	defer h.unlock()

	if r.URL.RequestURI() == healthPath {
		rw.WriteHeader(200)
		rw.Write(imgproxyIsRunningMsg)
		return
	}

	ctx, timeoutCancel := startTimer(ctx, time.Duration(conf.WriteTimeout)*time.Second)
	defer timeoutCancel()

	ctx, err := parsePath(ctx, r)
	if err != nil {
		panic(newError(404, err.Error(), "Invalid image url"))
	}

	ctx, downloadcancel, err := downloadImage(ctx)
	defer downloadcancel()
	if err != nil {
		if newRelicEnabled {
			sendErrorToNewRelic(ctx, err)
		}
		panic(newError(404, err.Error(), "Image is unreachable"))
	}

	checkTimeout(ctx)

	if conf.ETagEnabled {
		eTag := calcETag(ctx)
		rw.Header().Set("ETag", eTag)

		if eTag == r.Header.Get("If-None-Match") {
			panic(errNotModified)
		}
	}

	checkTimeout(ctx)

	imageData, err := processImage(ctx)
	if err != nil {
		if newRelicEnabled {
			sendErrorToNewRelic(ctx, err)
		}
		panic(newError(500, err.Error(), "Error occurred while processing image"))
	}

	checkTimeout(ctx)

	respondWithImage(ctx, reqID, r, rw, imageData)
}
<|endoftext|>


<commit_before>// Copyright 2020 Google LLC
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// 	https://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package xsrf

import (
	"context"
	"errors"
	"github.com/google/go-safeweb/safehttp"
	"golang.org/x/net/xsrftoken"
)

const (
	// TokenKey is the form key used when sending the token as part of POST
	// request.
	TokenKey = "xsrf-token"
)

var statePreservingMethods = map[string]bool{
	safehttp.MethodGet:     true,
	safehttp.MethodHead:    true,
	safehttp.MethodOptions: true,
}

// UserIdentifier provides the web application users' identifiers,
// needed in generating the XSRF token.
type UserIdentifier interface {
	// UserID returns the user's identifier based on the
	// safehttp.IncomingRequest received.
	UserID(*safehttp.IncomingRequest) (string, error)
}

// Interceptor implements XSRF protection.
type Interceptor struct {
	// AppKey uniquely identifies each registered service and should have high
	// entropy as it is use in generating the XSRF token.
	AppKey string
	// Identifier supports retrieving the user ID of application's user based on
	// incoming requests. This is needed in generating the XSRF token.
	Identifier UserIdentifier
}

type tokenCtxKey struct{}

// Token extracts the XSRF token from the incoming request. If it is not
// present, it returns a non-nil error.
func Token(r *safehttp.IncomingRequest) (string, error) {
	tok := r.Context().Value(tokenCtxKey{})
	if tok == nil {
		return "", errors.New("xsrf token not found")
	}
	return tok.(string), nil
}

// Before should be executed before directing the safehttp.IncomingRequest to
// the handler to ensure it is not part of the Cross-Site Request
// Forgery attack.
//
// In case of state changing requests (all except GET, HEAD and OPTIONS), it
// checks for the presence of an XSRF token in the request and validates it
// based on the user ID associated with the request.
//
// For authorized requests, it adds a cryptographically safe XSRF token to the
// incoming request. It can be later extracted using Token.
func (i *Interceptor) Before(w *safehttp.ResponseWriter, r *safehttp.IncomingRequest, cfg interface{}) safehttp.Result {
	userID, err := i.Identifier.UserID(r)
	if err != nil {
		return w.ClientError(safehttp.StatusUnauthorized)
	}

	actionID := r.Method() + " " + r.URL.Path()
	needsValidation := !statePreservingMethods[r.Method()]
	if needsValidation {
		f, err := r.PostForm()
		if err != nil {
			mf, err := r.MultipartForm(32 << 20)
			if err != nil {
				return w.ClientError(safehttp.StatusBadRequest)
			}
			f = &mf.Form
		}

		tok := f.String(TokenKey, "")
		if f.Err() != nil || tok == "" {
			return w.ClientError(safehttp.StatusUnauthorized)
		}

		if ok := xsrftoken.Valid(tok, i.AppKey, userID, actionID); !ok {
			return w.ClientError(safehttp.StatusForbidden)
		}
	}

	tok := xsrftoken.Generate(i.AppKey, userID, actionID)
	r.SetContext(context.WithValue(r.Context(), tokenCtxKey{}, tok))

	return safehttp.Result{}
}
<commit_msg>Small modification to godoc and addded a comment to explain why we fallback to multipart forms when the body doesn't contain a post formMade safehttp.ResponseWriter a pointer in Before.<commit_after>// Copyright 2020 Google LLC
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// 	https://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package xsrf

import (
	"context"
	"errors"
	"github.com/google/go-safeweb/safehttp"
	"golang.org/x/net/xsrftoken"
)

const (
	// TokenKey is the form key used when sending the token as part of POST
	// request.
	TokenKey = "xsrf-token"
)

var statePreservingMethods = map[string]bool{
	safehttp.MethodGet:     true,
	safehttp.MethodHead:    true,
	safehttp.MethodOptions: true,
}

// UserIdentifier provides the web application users' identifiers,
// needed in generating the XSRF token.
type UserIdentifier interface {
	// UserID returns the user's identifier based on the
	// safehttp.IncomingRequest received.
	UserID(*safehttp.IncomingRequest) (string, error)
}

// Interceptor implements XSRF protection.
type Interceptor struct {
	// AppKey uniquely identifies each registered service and should have high
	// entropy as it is used for generating the XSRF token.
	AppKey string
	// Identifier supports retrieving the user ID based on the incoming
	// request. This is needed for generating the XSRF token.
	Identifier UserIdentifier
}

type tokenCtxKey struct{}

// Token extracts the XSRF token from the incoming request. If it is not
// present, it returns a non-nil error.
func Token(r *safehttp.IncomingRequest) (string, error) {
	tok := r.Context().Value(tokenCtxKey{})
	if tok == nil {
		return "", errors.New("xsrf token not found")
	}
	return tok.(string), nil
}

// Before should be executed before directing the safehttp.IncomingRequest to
// the handler to ensure it is not part of the Cross-Site Request
// Forgery attack.
//
// In case of state changing requests (all except GET, HEAD and OPTIONS), it
// checks for the presence of an XSRF token in the request and validates it
// based on the user ID associated with the request.
//
// For authorized requests, it adds a cryptographically safe XSRF token to the
// incoming request. It can be later extracted using Token.
func (i *Interceptor) Before(w *safehttp.ResponseWriter, r *safehttp.IncomingRequest, cfg interface{}) safehttp.Result {
	userID, err := i.Identifier.UserID(r)
	if err != nil {
		return w.ClientError(safehttp.StatusUnauthorized)
	}

	actionID := r.Method() + " " + r.URL.Path()
	needsValidation := !statePreservingMethods[r.Method()]
	if needsValidation {
		f, err := r.PostForm()
		if err != nil {
			// We fallback to checking whether the form is multipart. Both types
			// are valid in an incoming request as long as the XSRF token is
			// present.
			mf, err := r.MultipartForm(32 << 20)
			if err != nil {
				return w.ClientError(safehttp.StatusBadRequest)
			}
			f = &mf.Form
		}

		tok := f.String(TokenKey, "")
		if f.Err() != nil || tok == "" {
			return w.ClientError(safehttp.StatusUnauthorized)
		}

		if ok := xsrftoken.Valid(tok, i.AppKey, userID, actionID); !ok {
			return w.ClientError(safehttp.StatusForbidden)
		}
	}

	tok := xsrftoken.Generate(i.AppKey, userID, actionID)
	r.SetContext(context.WithValue(r.Context(), tokenCtxKey{}, tok))

	return safehttp.Result{}
}
<|endoftext|>


<commit_before>package main

import (
	"bytes"
	"context"
	"crypto/subtle"
	"fmt"
	"log"
	"net/http"
	"sync"
	"time"

	nanoid "github.com/matoous/go-nanoid"
	"github.com/valyala/fasthttp"
)

var (
	mimes = map[imageType]string{
		imageTypeJPEG: "image/jpeg",
		imageTypePNG:  "image/png",
		imageTypeWEBP: "image/webp",
	}

	responseBufPool = sync.Pool{
		New: func() interface{} {
			return new(bytes.Buffer)
		},
	}

	authHeaderMust []byte

	healthRequestURI = []byte("/health")

	serverMutex mutex
)

func startServer() *fasthttp.Server {
	serverMutex = newMutex(conf.Concurrency)

	s := &fasthttp.Server{
		Name:        "imgproxy",
		Handler:     serveHTTP,
		Concurrency: conf.MaxClients,
		ReadTimeout: time.Duration(conf.ReadTimeout) * time.Second,
	}

	go func() {
		log.Printf("Starting server at %s\n", conf.Bind)
		if err := s.ListenAndServe(conf.Bind); err != nil {
			log.Fatalln(err)
		}
	}()

	return s
}

func shutdownServer(s *fasthttp.Server) {
	log.Println("Shutting down the server...")
	s.Shutdown()
}

func logResponse(status int, msg string) {
	var color int

	if status >= 500 {
		color = 31
	} else if status >= 400 {
		color = 33
	} else {
		color = 32
	}

	log.Printf("|\033[7;%dm %d \033[0m| %s\n", color, status, msg)
}

func writeCORS(rctx *fasthttp.RequestCtx) {
	if len(conf.AllowOrigin) > 0 {
		rctx.Request.Header.Set("Access-Control-Allow-Origin", conf.AllowOrigin)
		rctx.Request.Header.Set("Access-Control-Allow-Methods", "GET, OPTIONs")
	}
}

func respondWithImage(ctx context.Context, reqID string, rctx *fasthttp.RequestCtx, data []byte) {
	rctx.SetStatusCode(200)

	po := getprocessingOptions(ctx)

	rctx.SetContentType(mimes[po.Format])
	rctx.Response.Header.Set("Cache-Control", fmt.Sprintf("max-age=%d, public", conf.TTL))
	rctx.Response.Header.Set("Expires", time.Now().Add(time.Second*time.Duration(conf.TTL)).Format(http.TimeFormat))

	if conf.GZipCompression > 0 && rctx.Request.Header.HasAcceptEncodingBytes([]byte("gzip")) {
		rctx.Response.Header.Set("Content-Encoding", "gzip")
		gzipData(data, rctx)
	} else {
		rctx.SetBody(data)
	}

	logResponse(200, fmt.Sprintf("[%s] Processed in %s: %s; %+v", reqID, getTimerSince(ctx), getImageURL(ctx), po))
}

func respondWithError(reqID string, rctx *fasthttp.RequestCtx, err imgproxyError) {
	logResponse(err.StatusCode, fmt.Sprintf("[%s] %s", reqID, err.Message))

	rctx.SetStatusCode(err.StatusCode)
	rctx.SetBodyString(err.PublicMessage)
}

func respondWithOptions(reqID string, rctx *fasthttp.RequestCtx) {
	logResponse(200, fmt.Sprintf("[%s] Respond with options", reqID))
	rctx.SetStatusCode(200)
}

func prepareAuthHeaderMust() []byte {
	if len(authHeaderMust) == 0 {
		buf := bytes.NewBufferString("Bearer ")
		buf.WriteString(conf.Secret)
		authHeaderMust = []byte(conf.Secret)
	}

	return authHeaderMust
}

func checkSecret(rctx *fasthttp.RequestCtx) bool {
	if len(conf.Secret) == 0 {
		return true
	}

	return subtle.ConstantTimeCompare(
		rctx.Request.Header.Peek("Authorization"),
		prepareAuthHeaderMust(),
	) == 1
}

func serveHTTP(rctx *fasthttp.RequestCtx) {
	reqID, _ := nanoid.Nanoid()

	defer func() {
		if r := recover(); r != nil {
			if err, ok := r.(imgproxyError); ok {
				respondWithError(reqID, rctx, err)
			} else {
				respondWithError(reqID, rctx, newUnexpectedError(r.(error), 4))
			}
		}
	}()

	log.Printf("[%s] %s: %s\n", reqID, rctx.Method(), rctx.RequestURI())

	writeCORS(rctx)

	if rctx.Request.Header.IsOptions() {
		respondWithOptions(reqID, rctx)
		return
	}

	if !rctx.IsGet() {
		panic(invalidMethodErr)
	}

	if !checkSecret(rctx) {
		panic(invalidSecretErr)
	}

	serverMutex.Lock()
	defer serverMutex.Unock()

	if bytes.Equal(rctx.RequestURI(), healthRequestURI) {
		rctx.SetStatusCode(200)
		rctx.SetBodyString("imgproxy is running")
		return
	}

	ctx, timeoutCancel := startTimer(time.Duration(conf.WriteTimeout) * time.Second)
	defer timeoutCancel()

	ctx, err := parsePath(ctx, rctx)
	if err != nil {
		panic(newError(404, err.Error(), "Invalid image url"))
	}

	ctx, downloadcancel, err := downloadImage(ctx)
	defer downloadcancel()
	if err != nil {
		panic(newError(404, err.Error(), "Image is unreachable"))
	}

	checkTimeout(ctx)

	// if conf.ETagEnabled {
	// 	eTag := calcETag(b, &procOpt)
	// 	rw.Header().Set("ETag", eTag)

	// 	if eTag == r.Header.Get("If-None-Match") {
	// 		panic(notModifiedErr)
	// 	}
	// }

	checkTimeout(ctx)

	imageData, err := processImage(ctx)
	if err != nil {
		panic(newError(500, err.Error(), "Error occurred while processing image"))
	}

	checkTimeout(ctx)

	respondWithImage(ctx, reqID, rctx, imageData)
}
<commit_msg>Fix secret checking<commit_after>package main

import (
	"bytes"
	"context"
	"crypto/subtle"
	"fmt"
	"log"
	"net/http"
	"sync"
	"time"

	nanoid "github.com/matoous/go-nanoid"
	"github.com/valyala/fasthttp"
)

var (
	mimes = map[imageType]string{
		imageTypeJPEG: "image/jpeg",
		imageTypePNG:  "image/png",
		imageTypeWEBP: "image/webp",
	}

	responseBufPool = sync.Pool{
		New: func() interface{} {
			return new(bytes.Buffer)
		},
	}

	authHeaderMust []byte

	healthRequestURI = []byte("/health")

	serverMutex mutex
)

func startServer() *fasthttp.Server {
	serverMutex = newMutex(conf.Concurrency)

	s := &fasthttp.Server{
		Name:        "imgproxy",
		Handler:     serveHTTP,
		Concurrency: conf.MaxClients,
		ReadTimeout: time.Duration(conf.ReadTimeout) * time.Second,
	}

	go func() {
		log.Printf("Starting server at %s\n", conf.Bind)
		if err := s.ListenAndServe(conf.Bind); err != nil {
			log.Fatalln(err)
		}
	}()

	return s
}

func shutdownServer(s *fasthttp.Server) {
	log.Println("Shutting down the server...")
	s.Shutdown()
}

func logResponse(status int, msg string) {
	var color int

	if status >= 500 {
		color = 31
	} else if status >= 400 {
		color = 33
	} else {
		color = 32
	}

	log.Printf("|\033[7;%dm %d \033[0m| %s\n", color, status, msg)
}

func writeCORS(rctx *fasthttp.RequestCtx) {
	if len(conf.AllowOrigin) > 0 {
		rctx.Request.Header.Set("Access-Control-Allow-Origin", conf.AllowOrigin)
		rctx.Request.Header.Set("Access-Control-Allow-Methods", "GET, OPTIONs")
	}
}

func respondWithImage(ctx context.Context, reqID string, rctx *fasthttp.RequestCtx, data []byte) {
	rctx.SetStatusCode(200)

	po := getprocessingOptions(ctx)

	rctx.SetContentType(mimes[po.Format])
	rctx.Response.Header.Set("Cache-Control", fmt.Sprintf("max-age=%d, public", conf.TTL))
	rctx.Response.Header.Set("Expires", time.Now().Add(time.Second*time.Duration(conf.TTL)).Format(http.TimeFormat))

	if conf.GZipCompression > 0 && rctx.Request.Header.HasAcceptEncodingBytes([]byte("gzip")) {
		rctx.Response.Header.Set("Content-Encoding", "gzip")
		gzipData(data, rctx)
	} else {
		rctx.SetBody(data)
	}

	logResponse(200, fmt.Sprintf("[%s] Processed in %s: %s; %+v", reqID, getTimerSince(ctx), getImageURL(ctx), po))
}

func respondWithError(reqID string, rctx *fasthttp.RequestCtx, err imgproxyError) {
	logResponse(err.StatusCode, fmt.Sprintf("[%s] %s", reqID, err.Message))

	rctx.SetStatusCode(err.StatusCode)
	rctx.SetBodyString(err.PublicMessage)
}

func respondWithOptions(reqID string, rctx *fasthttp.RequestCtx) {
	logResponse(200, fmt.Sprintf("[%s] Respond with options", reqID))
	rctx.SetStatusCode(200)
}

func prepareAuthHeaderMust() []byte {
	if len(authHeaderMust) == 0 {
		authHeaderMust = []byte(fmt.Sprintf("Bearer %s", conf.Secret))
	}

	return authHeaderMust
}

func checkSecret(rctx *fasthttp.RequestCtx) bool {
	if len(conf.Secret) == 0 {
		return true
	}

	return subtle.ConstantTimeCompare(
		rctx.Request.Header.Peek("Authorization"),
		prepareAuthHeaderMust(),
	) == 1
}

func serveHTTP(rctx *fasthttp.RequestCtx) {
	reqID, _ := nanoid.Nanoid()

	defer func() {
		if r := recover(); r != nil {
			if err, ok := r.(imgproxyError); ok {
				respondWithError(reqID, rctx, err)
			} else {
				respondWithError(reqID, rctx, newUnexpectedError(r.(error), 4))
			}
		}
	}()

	log.Printf("[%s] %s: %s\n", reqID, rctx.Method(), rctx.RequestURI())

	writeCORS(rctx)

	if rctx.Request.Header.IsOptions() {
		respondWithOptions(reqID, rctx)
		return
	}

	if !rctx.IsGet() {
		panic(invalidMethodErr)
	}

	if !checkSecret(rctx) {
		panic(invalidSecretErr)
	}

	serverMutex.Lock()
	defer serverMutex.Unock()

	if bytes.Equal(rctx.RequestURI(), healthRequestURI) {
		rctx.SetStatusCode(200)
		rctx.SetBodyString("imgproxy is running")
		return
	}

	ctx, timeoutCancel := startTimer(time.Duration(conf.WriteTimeout) * time.Second)
	defer timeoutCancel()

	ctx, err := parsePath(ctx, rctx)
	if err != nil {
		panic(newError(404, err.Error(), "Invalid image url"))
	}

	ctx, downloadcancel, err := downloadImage(ctx)
	defer downloadcancel()
	if err != nil {
		panic(newError(404, err.Error(), "Image is unreachable"))
	}

	checkTimeout(ctx)

	// if conf.ETagEnabled {
	// 	eTag := calcETag(b, &procOpt)
	// 	rw.Header().Set("ETag", eTag)

	// 	if eTag == r.Header.Get("If-None-Match") {
	// 		panic(notModifiedErr)
	// 	}
	// }

	checkTimeout(ctx)

	imageData, err := processImage(ctx)
	if err != nil {
		panic(newError(500, err.Error(), "Error occurred while processing image"))
	}

	checkTimeout(ctx)

	respondWithImage(ctx, reqID, rctx, imageData)
}
<|endoftext|>