<commit_before>package main

import (
	"crypto/sha1"
	"encoding/hex"
	"errors"
	"flag"
	"fmt"
	"github.com/bmizerany/pat"
	"github.com/fzzbt/radix/redis"
	"io"
	"io/ioutil"
	"log"
	"net/http"
	"os"
	"path"
	"strconv"
	"strings"
	"syscall"
)

// The maximal size for an image is 5MB
const maxSize = 5 * (1 << 20)

// The directory for caching files
var directory string

// The connection to redis
var connection *redis.Client

func urlStatus(uri string) error {
	ok, err := connection.Hexists("img/"+uri, "created_at").Bool()
	if err != nil {
		return err
	}
	if !ok {
		return errors.New("Invalid URL")
	}

	str, err := connection.Get("img/err/" + uri).Str()
	if err == nil {
		return errors.New(str)
	}

	return nil
}

// Generate a key for cache from a string
func generateKeyForCache(s string) string {
	h := sha1.New()
	io.WriteString(h, s)
	key := h.Sum(nil)

	// Use 3 levels of hasing to avoid having too many files in the same directory
	return fmt.Sprintf("%s/%x/%x/%x/%x", directory, key[0:1], key[1:2], key[2:3], key[3:])
}

// Fetch image from cache
func fetchImageFromCache(uri string) (contentType string, body []byte, ok bool) {
	ok = false

	contentType, err := connection.Hget("img/"+uri, "type").Str()
	if err != nil {
		return
	}

	filename := generateKeyForCache(uri)
	body, err = ioutil.ReadFile(filename)
	ok = err == nil

	return
}

// Save the body and the content-type header in cache
func saveImageInCache(uri string, contentType string, body []byte) {
	go func() {
		filename := generateKeyForCache(uri)
		dirname := path.Dir(filename)
		err := os.MkdirAll(dirname, 0755)
		if err != nil {
			return
		}

		// Save the body on disk
		err = ioutil.WriteFile(filename, body, 0644)
		if err != nil {
			log.Printf("Error while writing %s\n", filename)
			return
		}

		// And other infos in redis
		connection.Hset("img/"+uri, "type", contentType)
	}()
}

func saveErrorInCache(uri string, err error) {
	go func() {
		connection.Set("img/err/"+uri, err.Error())
		connection.Expire("img/err/"+uri, 600)
	}()
}

// Fetch the image from the distant server
func fetchImageFromServer(uri string) (contentType string, body []byte, err error) {
	res, err := http.Get(uri)
	if err != nil {
		return
	}
	if res.StatusCode != 200 {
		log.Printf("Status code of %s is: %d\n", uri, res.StatusCode)
		err = errors.New("Unexpected status code")
		saveErrorInCache(uri, err)
		return
	}

	defer res.Body.Close()
	body, err = ioutil.ReadAll(res.Body)
	if err != nil {
		return
	}
	if res.ContentLength > maxSize {
		log.Printf("Exceeded max size for %s: %d\n", uri, res.ContentLength)
		err = errors.New("Exceeded max size")
		saveErrorInCache(uri, err)
		return
	}
	contentType = res.Header.Get("Content-Type")
	if contentType[0:5] != "image" {
		log.Printf("%s has an invalid content-type: %s\n", uri, contentType)
		err = errors.New("Invalid content-type")
		saveErrorInCache(uri, err)
		return
	}
	log.Printf("Fetch %s (%s)\n", uri, contentType)

	if urlStatus(uri) == nil {
		saveImageInCache(uri, contentType, body)
	}
	return
}

// Fetch image from cache if available, or from the server
func fetchImage(uri string) (contentType string, body []byte, err error) {
	err = urlStatus(uri)
	if err != nil {
		return
	}

	contentType, body, ok := fetchImageFromCache(uri)
	if !ok {
		contentType, body, err = fetchImageFromServer(uri)
	}

	return
}

// Receive an HTTP request for an image and respond with it
func Img(w http.ResponseWriter, r *http.Request) {
	encoded_url := r.URL.Query().Get(":encoded_url")
	chars, err := hex.DecodeString(encoded_url)
	if err != nil {
		log.Printf("Invalid URL %s\n", encoded_url)
		http.Error(w, "Invalid parameters", 400)
		return
	}
	uri := string(chars)

	contentType, body, err := fetchImage(uri)
	if err != nil {
		http.NotFound(w, r)
		return
	}
	w.Header().Add("Content-Type", contentType)
	w.Write(body)
}

// Returns 200 OK if the server is running (for monitoring)
func Status(w http.ResponseWriter, r *http.Request) {
	fmt.Fprintf(w, "OK")
}

func main() {
	// Parse the command-line
	var addr string
	var logs string
	var conn string
	flag.StringVar(&addr, "a", "127.0.0.1:8000", "Bind to this address:port")
	flag.StringVar(&logs, "l", "-", "Use this file for logs")
	flag.StringVar(&conn, "r", "localhost:6379/0", "The redis database to use for caching meta")
	flag.StringVar(&directory, "d", "cache", "The directory for the caching files")
	flag.Parse()
<commit_msg>Add HTTP headers for caching<commit_after>package main

import (
	"crypto/sha1"
	"encoding/hex"
	"errors"
	"flag"
	"fmt"
	"github.com/bmizerany/pat"
	"github.com/fzzbt/radix/redis"
	"io"
	"io/ioutil"
	"log"
	"net/http"
	"os"
	"path"
	"strconv"
	"strings"
	"syscall"
	"time"
)

// HTTP headers struct
type Headers struct {
	contentType  string
	lastModified string
	cacheControl string
}

// The maximal size for an image is 5MB
const maxSize = 5 * (1 << 20)

// The directory for caching files
var directory string

// The connection to redis
var connection *redis.Client

func urlStatus(uri string) error {
	ok, err := connection.Hexists("img/"+uri, "created_at").Bool()
	if err != nil {
		return err
	}
	if !ok {
		return errors.New("Invalid URL")
	}

	str, err := connection.Get("img/err/" + uri).Str()
	if err == nil {
		return errors.New(str)
	}

	return nil
}

// Generate a key for cache from a string
func generateKeyForCache(s string) string {
	h := sha1.New()
	io.WriteString(h, s)
	key := h.Sum(nil)

	// Use 3 levels of hasing to avoid having too many files in the same directory
	return fmt.Sprintf("%s/%x/%x/%x/%x", directory, key[0:1], key[1:2], key[2:3], key[3:])
}

// Fetch image from cache
func fetchImageFromCache(uri string) (headers Headers, body []byte, ok bool) {
	ok = false

	contentType, err := connection.Hget("img/"+uri, "type").Str()
	if err != nil {
		return
	}

	filename := generateKeyForCache(uri)
	stat, err := os.Stat(filename)
	if err != nil {
		return
	}

	headers.contentType = contentType
	headers.lastModified = stat.ModTime().Format(time.RFC1123)

	body, err = ioutil.ReadFile(filename)
	ok = err == nil

	return
}

// Save the body and the content-type header in cache
func saveImageInCache(uri string, headers Headers, body []byte) {
	go func() {
		filename := generateKeyForCache(uri)
		dirname := path.Dir(filename)
		err := os.MkdirAll(dirname, 0755)
		if err != nil {
			return
		}

		// Save the body on disk
		err = ioutil.WriteFile(filename, body, 0644)
		if err != nil {
			log.Printf("Error while writing %s\n", filename)
			return
		}

		// And other infos in redis
		connection.Hset("img/"+uri, "type", headers.contentType)
	}()
}

func saveErrorInCache(uri string, err error) {
	go func() {
		connection.Set("img/err/"+uri, err.Error())
		connection.Expire("img/err/"+uri, 600)
	}()
}

// Fetch the image from the distant server
func fetchImageFromServer(uri string) (headers Headers, body []byte, err error) {
	res, err := http.Get(uri)
	if err != nil {
		return
	}
	if res.StatusCode != 200 {
		log.Printf("Status code of %s is: %d\n", uri, res.StatusCode)
		err = errors.New("Unexpected status code")
		saveErrorInCache(uri, err)
		return
	}

	defer res.Body.Close()
	body, err = ioutil.ReadAll(res.Body)
	if err != nil {
		return
	}
	if res.ContentLength > maxSize {
		log.Printf("Exceeded max size for %s: %d\n", uri, res.ContentLength)
		err = errors.New("Exceeded max size")
		saveErrorInCache(uri, err)
		return
	}
	contentType := res.Header.Get("Content-Type")
	if contentType[0:5] != "image" {
		log.Printf("%s has an invalid content-type: %s\n", uri, contentType)
		err = errors.New("Invalid content-type")
		saveErrorInCache(uri, err)
		return
	}
	log.Printf("Fetch %s (%s)\n", uri, contentType)

	headers.contentType = contentType
	headers.lastModified = time.Now().Format(time.RFC1123)
	if urlStatus(uri) == nil {
		saveImageInCache(uri, headers, body)
	}
	return
}

// Fetch image from cache if available, or from the server
func fetchImage(uri string) (headers Headers, body []byte, err error) {
	err = urlStatus(uri)
	if err != nil {
		return
	}

	headers, body, ok := fetchImageFromCache(uri)
	if !ok {
		headers, body, err = fetchImageFromServer(uri)
	}
	headers.cacheControl = "public, max-age=600"

	return
}

// Receive an HTTP request for an image and respond with it
func Img(w http.ResponseWriter, r *http.Request) {
	encoded_url := r.URL.Query().Get(":encoded_url")
	chars, err := hex.DecodeString(encoded_url)
	if err != nil {
		log.Printf("Invalid URL %s\n", encoded_url)
		http.Error(w, "Invalid parameters", 400)
		return
	}
	uri := string(chars)

	headers, body, err := fetchImage(uri)
	if err != nil {
		http.NotFound(w, r)
		return
	}
	if headers.lastModified == r.Header.Get("If-Modified-Since") {
		w.WriteHeader(http.StatusNotModified)
		return
	}
	w.Header().Add("Content-Type", headers.contentType)
	w.Header().Add("Last-Modified", headers.lastModified)
	w.Header().Add("Cache-Control", headers.cacheControl)
	w.Write(body)
}

// Returns 200 OK if the server is running (for monitoring)
func Status(w http.ResponseWriter, r *http.Request) {
	fmt.Fprintf(w, "OK")
}

func main() {
	// Parse the command-line
	var addr string
	var logs string
	var conn string
	flag.StringVar(&addr, "a", "127.0.0.1:8000", "Bind to this address:port")
	flag.StringVar(&logs, "l", "-", "Use this file for logs")
	flag.StringVar(&conn, "r", "localhost:6379/0", "The redis database to use for caching meta")
	flag.StringVar(&directory, "d", "cache", "The directory for the caching files")
	flag.Parse()

{'fraction_non_alphanumeric': 0.14822217546812538, 'fraction_numerical': 0.01209762255417631, 'mean_word_length': 3.9055727554179565, 'pattern_counts': {'":': 2, '<': 7, '<?xml version=': 0, '>': 5, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 2, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '1096674', 'n_tokens_mistral': 3596, 'n_tokens_neox': 3409, 'n_words': 970}