<commit_before>package proxyproto

import (
	"bufio"
	"net"
	"sync"
	"time"
)

// Listener is used to wrap an underlying listener,
// whose connections may be using the HAProxy Proxy Protocol.
// If the connection is using the protocol, the RemoteAddr() will return
// the correct client address.
//
// Optionally define ProxyHeaderTimeout to set a maximum time to
// receive the Proxy Protocol Header. Zero means no timeout.
type Listener struct {
	Listener           net.Listener
	ProxyHeaderTimeout time.Duration
}

// Conn is used to wrap and underlying connection which
// may be speaking the Proxy Protocol. If it is, the RemoteAddr() will
// return the address of the client instead of the proxy address.
type Conn struct {
	bufReader          *bufio.Reader
	conn               net.Conn
	header             *Header
	once               sync.Once
	proxyHeaderTimeout time.Duration
}

// Accept waits for and returns the next connection to the listener.
func (p *Listener) Accept() (net.Conn, error) {
	// Get the underlying connection
	conn, err := p.Listener.Accept()
	if err != nil {
		return nil, err
	}
	return NewConn(conn, p.ProxyHeaderTimeout), nil
}

// Close closes the underlying listener.
func (p *Listener) Close() error {
	return p.Listener.Close()
}

// Addr returns the underlying listener's network address.
func (p *Listener) Addr() net.Addr {
	return p.Listener.Addr()
}

// NewConn is used to wrap a net.Conn that may be speaking
// the proxy protocol into a proxyproto.Conn
func NewConn(conn net.Conn, timeout time.Duration) *Conn {
	pConn := &Conn{
		bufReader:          bufio.NewReader(conn),
		conn:               conn,
		proxyHeaderTimeout: timeout,
	}
	return pConn
}

// Read is check for the proxy protocol header when doing
// the initial scan. If there is an error parsing the header,
// it is returned and the socket is closed.
func (p *Conn) Read(b []byte) (int, error) {
	var err error
	p.once.Do(func() {
		err = p.readHeader()
	})
	if err != nil {
		return 0, err
	}
	return p.bufReader.Read(b)
}

// Write wraps original conn.Write
func (p *Conn) Write(b []byte) (int, error) {
	return p.conn.Write(b)
}

// Close wraps original conn.Close
func (p *Conn) Close() error {
	return p.conn.Close()
}

// LocalAddr returns the address of the server if the proxy
// protocol is being used, otherwise just returns the address of
// the socket server.
func (p *Conn) LocalAddr() net.Addr {
	p.once.Do(func() { p.readHeader() })
	if p.header == nil {
		return p.conn.LocalAddr()
	}

	return p.header.LocalAddr()
}

// RemoteAddr returns the address of the client if the proxy
// protocol is being used, otherwise just returns the address of
// the socket peer.
func (p *Conn) RemoteAddr() net.Addr {
	p.once.Do(func() { p.readHeader() })
	if p.header == nil {
		return p.conn.RemoteAddr()
	}

	return p.header.RemoteAddr()
}

// SetDeadline wraps original conn.SetDeadline
func (p *Conn) SetDeadline(t time.Time) error {
	return p.conn.SetDeadline(t)
}

// SetReadDeadline wraps original conn.SetReadDeadline
func (p *Conn) SetReadDeadline(t time.Time) error {
	return p.conn.SetReadDeadline(t)
}

// SetWriteDeadline wraps original conn.SetWriteDeadline
func (p *Conn) SetWriteDeadline(t time.Time) error {
	return p.conn.SetWriteDeadline(t)
}

func (p *Conn) readHeader() error {
	p.header, _ = Read(p.bufReader)
	return nil

}
<commit_msg>protocol: don't swallow connection errors<commit_after>package proxyproto

import (
	"bufio"
	"net"
	"sync"
	"time"
)

// Listener is used to wrap an underlying listener,
// whose connections may be using the HAProxy Proxy Protocol.
// If the connection is using the protocol, the RemoteAddr() will return
// the correct client address.
//
// Optionally define ProxyHeaderTimeout to set a maximum time to
// receive the Proxy Protocol Header. Zero means no timeout.
type Listener struct {
	Listener           net.Listener
	ProxyHeaderTimeout time.Duration
}

// Conn is used to wrap and underlying connection which
// may be speaking the Proxy Protocol. If it is, the RemoteAddr() will
// return the address of the client instead of the proxy address.
type Conn struct {
	bufReader          *bufio.Reader
	conn               net.Conn
	header             *Header
	once               sync.Once
	proxyHeaderTimeout time.Duration
}

// Accept waits for and returns the next connection to the listener.
func (p *Listener) Accept() (net.Conn, error) {
	// Get the underlying connection
	conn, err := p.Listener.Accept()
	if err != nil {
		return nil, err
	}
	return NewConn(conn, p.ProxyHeaderTimeout), nil
}

// Close closes the underlying listener.
func (p *Listener) Close() error {
	return p.Listener.Close()
}

// Addr returns the underlying listener's network address.
func (p *Listener) Addr() net.Addr {
	return p.Listener.Addr()
}

// NewConn is used to wrap a net.Conn that may be speaking
// the proxy protocol into a proxyproto.Conn
func NewConn(conn net.Conn, timeout time.Duration) *Conn {
	pConn := &Conn{
		bufReader:          bufio.NewReader(conn),
		conn:               conn,
		proxyHeaderTimeout: timeout,
	}
	return pConn
}

// Read is check for the proxy protocol header when doing
// the initial scan. If there is an error parsing the header,
// it is returned and the socket is closed.
func (p *Conn) Read(b []byte) (int, error) {
	var err error
	p.once.Do(func() {
		err = p.readHeader()
	})
	if err != nil {
		return 0, err
	}
	return p.bufReader.Read(b)
}

// Write wraps original conn.Write
func (p *Conn) Write(b []byte) (int, error) {
	return p.conn.Write(b)
}

// Close wraps original conn.Close
func (p *Conn) Close() error {
	return p.conn.Close()
}

// LocalAddr returns the address of the server if the proxy
// protocol is being used, otherwise just returns the address of
// the socket server.
func (p *Conn) LocalAddr() net.Addr {
	p.once.Do(func() { p.readHeader() })
	if p.header == nil {
		return p.conn.LocalAddr()
	}

	return p.header.LocalAddr()
}

// RemoteAddr returns the address of the client if the proxy
// protocol is being used, otherwise just returns the address of
// the socket peer.
func (p *Conn) RemoteAddr() net.Addr {
	p.once.Do(func() { p.readHeader() })
	if p.header == nil {
		return p.conn.RemoteAddr()
	}

	return p.header.RemoteAddr()
}

// SetDeadline wraps original conn.SetDeadline
func (p *Conn) SetDeadline(t time.Time) error {
	return p.conn.SetDeadline(t)
}

// SetReadDeadline wraps original conn.SetReadDeadline
func (p *Conn) SetReadDeadline(t time.Time) error {
	return p.conn.SetReadDeadline(t)
}

// SetWriteDeadline wraps original conn.SetWriteDeadline
func (p *Conn) SetWriteDeadline(t time.Time) error {
	return p.conn.SetWriteDeadline(t)
}

func (p *Conn) readHeader() (err error) {
	p.header, err = Read(p.bufReader)
	// For the purpose of this wrapper shamefully stolen from armon/go-proxyproto
	// let's act as if there was no error when PROXY protocol is not present.
	if err == ErrNoProxyProtocol {
		err = nil
	}

	return
}
<|endoftext|>


<commit_before>package centrifuge

import (
	"encoding/json"
)

type clientCommand struct {
	UID    string `json:"uid"`
	Method string `json:"method"`
}

type connectClientCommand struct {
	clientCommand
	Params connectParams `json:"params"`
}

type refreshClientCommand struct {
	clientCommand
	Params refreshParams `json:"params"`
}

type subscribeClientCommand struct {
	clientCommand
	Params subscribeParams `json:"params"`
}

type unsubscribeClientCommand struct {
	clientCommand
	Params unsubscribeParams `json:"params"`
}

type publishClientCommand struct {
	clientCommand
	Params publishParams `json:"params"`
}

type presenceClientCommand struct {
	clientCommand
	Params presenceParams `json:"params"`
}

type historyClientCommand struct {
	clientCommand
	Params historyParams `json:"params"`
}

type pingClientCommand struct {
	clientCommand
}

type connectParams struct {
	User      string `json:"user"`
	Timestamp string `json:"timestamp"`
	Info      string `json:"info"`
	Token     string `json:"token"`
}

type refreshParams struct {
	User      string `json:"user"`
	Timestamp string `json:"timestamp"`
	Info      string `json:"info"`
	Token     string `json:"token"`
}

type subscribeParams struct {
	Channel string `json:"channel"`
	Client  string `json:"client"`
	Last    string `json:"last"`
	Recover bool   `json:"recover"`
	Info    string `json:"info"`
	Sign    string `json:"sign"`
}

type unsubscribeParams struct {
	Channel string `json:"channel"`
}

type publishParams struct {
	Channel string           `json:"channel"`
	Data    *json.RawMessage `json:"data"`
}

type presenceParams struct {
	Channel string `json:"channel"`
}

type historyParams struct {
	Channel string `json:"channel"`
}

type response struct {
	UID    string          `json:"uid,omitempty"`
	Error  string          `json:"error"`
	Method string          `json:"method"`
	Body   json.RawMessage `json:"body"`
}

type joinLeaveMessage struct {
	Channel string        `json:"channel"`
	Data    rawClientInfo `json:"data"`
}

type connectResponseBody struct {
	Version string `json:"version"`
	Client  string `json:"client"`
	Expires bool   `json:"expires"`
	Expired bool   `json:"expired"`
	TTL     int64  `json:"ttl"`
}

type subscribeResponseBody struct {
	Channel   string       `json:"channel"`
	Status    bool         `json:"status"`
	Last      string       `json:"last"`
	Messages  []rawMessage `json:"messages"`
	Recovered bool         `json:"recovered"`
}

type unsubscribeResponseBody struct {
	Channel string `json:"channel"`
	Status  bool   `json:"status"`
}

type publishResponseBody struct {
	Channel string `json:"channel"`
	Status  bool   `json:"status"`
}

type presenceResponseBody struct {
	Channel string                   `json:"channel"`
	Data    map[string]rawClientInfo `json:"data"`
}

type historyResponseBody struct {
	Channel string       `json:"channel"`
	Data    []rawMessage `json:"data"`
}

type disconnectAdvice struct {
	Reason    string `json:"reason"`
	Reconnect bool   `json:"reconnect"`
}

var (
	arrayJsonPrefix  byte = '['
	objectJsonPrefix byte = '{'
)

func responsesFromClientMsg(msg []byte) ([]response, error) {
	var resps []response
	firstByte := msg[0]
	switch firstByte {
	case objectJsonPrefix:
		// single command request
		var resp response
		err := json.Unmarshal(msg, &resp)
		if err != nil {
			return nil, err
		}
		resps = append(resps, resp)
	case arrayJsonPrefix:
		// array of commands received
		err := json.Unmarshal(msg, &resps)
		if err != nil {
			return nil, err
		}
	default:
		return nil, ErrInvalidMessage
	}
	return resps, nil
}
<commit_msg>omitempty for subscribe fields<commit_after>package centrifuge

import (
	"encoding/json"
)

type clientCommand struct {
	UID    string `json:"uid"`
	Method string `json:"method"`
}

type connectClientCommand struct {
	clientCommand
	Params connectParams `json:"params"`
}

type refreshClientCommand struct {
	clientCommand
	Params refreshParams `json:"params"`
}

type subscribeClientCommand struct {
	clientCommand
	Params subscribeParams `json:"params"`
}

type unsubscribeClientCommand struct {
	clientCommand
	Params unsubscribeParams `json:"params"`
}

type publishClientCommand struct {
	clientCommand
	Params publishParams `json:"params"`
}

type presenceClientCommand struct {
	clientCommand
	Params presenceParams `json:"params"`
}

type historyClientCommand struct {
	clientCommand
	Params historyParams `json:"params"`
}

type pingClientCommand struct {
	clientCommand
}

type connectParams struct {
	User      string `json:"user"`
	Timestamp string `json:"timestamp"`
	Info      string `json:"info"`
	Token     string `json:"token"`
}

type refreshParams struct {
	User      string `json:"user"`
	Timestamp string `json:"timestamp"`
	Info      string `json:"info"`
	Token     string `json:"token"`
}

type subscribeParams struct {
	Channel string `json:"channel"`
	Client  string `json:"client,omitempty"`
	Last    string `json:"last,omitempty"`
	Recover bool   `json:"recover,omitempty"`
	Info    string `json:"info,omitempty"`
	Sign    string `json:"sign,omitempty"`
}

type unsubscribeParams struct {
	Channel string `json:"channel"`
}

type publishParams struct {
	Channel string           `json:"channel"`
	Data    *json.RawMessage `json:"data"`
}

type presenceParams struct {
	Channel string `json:"channel"`
}

type historyParams struct {
	Channel string `json:"channel"`
}

type response struct {
	UID    string          `json:"uid,omitempty"`
	Error  string          `json:"error"`
	Method string          `json:"method"`
	Body   json.RawMessage `json:"body"`
}

type joinLeaveMessage struct {
	Channel string        `json:"channel"`
	Data    rawClientInfo `json:"data"`
}

type connectResponseBody struct {
	Version string `json:"version"`
	Client  string `json:"client"`
	Expires bool   `json:"expires"`
	Expired bool   `json:"expired"`
	TTL     int64  `json:"ttl"`
}

type subscribeResponseBody struct {
	Channel   string       `json:"channel"`
	Status    bool         `json:"status"`
	Last      string       `json:"last"`
	Messages  []rawMessage `json:"messages"`
	Recovered bool         `json:"recovered"`
}

type unsubscribeResponseBody struct {
	Channel string `json:"channel"`
	Status  bool   `json:"status"`
}

type publishResponseBody struct {
	Channel string `json:"channel"`
	Status  bool   `json:"status"`
}

type presenceResponseBody struct {
	Channel string                   `json:"channel"`
	Data    map[string]rawClientInfo `json:"data"`
}

type historyResponseBody struct {
	Channel string       `json:"channel"`
	Data    []rawMessage `json:"data"`
}

type disconnectAdvice struct {
	Reason    string `json:"reason"`
	Reconnect bool   `json:"reconnect"`
}

var (
	arrayJsonPrefix  byte = '['
	objectJsonPrefix byte = '{'
)

func responsesFromClientMsg(msg []byte) ([]response, error) {
	var resps []response
	firstByte := msg[0]
	switch firstByte {
	case objectJsonPrefix:
		// single command request
		var resp response
		err := json.Unmarshal(msg, &resp)
		if err != nil {
			return nil, err
		}
		resps = append(resps, resp)
	case arrayJsonPrefix:
		// array of commands received
		err := json.Unmarshal(msg, &resps)
		if err != nil {
			return nil, err
		}
	default:
		return nil, ErrInvalidMessage
	}
	return resps, nil
}
<|endoftext|>


<commit_before>package main

import (
	"encoding/json"
	"errors"
	"github.com/franela/goreq"
	"github.com/gin-gonic/gin"
	"log"
	"sync/atomic"
	"time"
)

var errTaskNotFresh = errors.New("This task has been running too long to request a token.")
var errAlreadyGivenKey = errors.New("This task has already been given a token.")
var usedTaskIds = NewTtlSet()

func createToken(token string, opts interface{}) (string, error) {
	r, err := goreq.Request{
		Uri:    vaultPath("/v1/auth/token/create", ""),
		Method: "POST",
		Body:   opts,
	}.WithHeader("X-Vault-Token", token).Do()
	if err == nil {
		defer r.Body.Close()
		switch r.StatusCode {
		case 200:
			var t vaultTokenResp
			if err := r.Body.FromJsonTo(&t); err == nil {
				return t.Auth.ClientToken, nil
			} else {
				return "", err
			}
		default:
			var e vaultError
			e.Code = r.StatusCode
			if err := r.Body.FromJsonTo(&e); err == nil {
				return "", e
			} else {
				e.Errors = []string{"communication error."}
				return "", e
			}
		}
	} else {
		return "", err
	}
}

func createTokenPair(token string, p *policy) (string, error) {
	tempTokenOpts := struct {
		Ttl      string   `json:"ttl"`
		NumUses  int      `json:"num_uses"`
		Policies []string `json:"policies"`
		NoParent bool     `json:"no_parent"`
	}{"10m", 2, []string{"default"}, true}
	pol := p.Policies
	if len(pol) == 0 { // explicitly set the policy, else the token will inherit ours
		pol = []string{"default"}
	}
	permTokenOpts := struct {
		Ttl      string            `json:"ttl,omitempty"`
		Policies []string          `json:"policies"`
		Meta     map[string]string `json:"meta,omitempty"`
		NumUses  int               `json:"num_uses"`
		NoParent bool              `json:"no_parent"`
	}{time.Duration(time.Duration(p.Ttl) * time.Second).String(), pol, p.Meta, p.NumUses, true}

	if tempToken, err := createToken(token, tempTokenOpts); err == nil {
		if permToken, err := createToken(token, permTokenOpts); err == nil {
			r, err := goreq.Request{
				Uri:    vaultPath("/v1/cubbyhole/vault-token", ""),
				Method: "POST",
				Body: struct {
					Token string `json:"token"`
				}{permToken},
			}.WithHeader("X-Vault-Token", tempToken).Do()
			if err == nil {
				defer r.Body.Close()
				switch r.StatusCode {
				case 204:
					return tempToken, nil
				default:
					var e vaultError
					e.Code = r.StatusCode
					if err := r.Body.FromJsonTo(&e); err == nil {
						return "", e
					} else {
						e.Errors = []string{"communication error."}
						return "", e
					}
				}
			} else {
				return "", err
			}
		} else {
			return "", err
		}
	} else {
		return "", err
	}
}

func Provide(c *gin.Context) {
	state.RLock()
	status := state.Status
	token := state.Token
	state.RUnlock()

	remoteIp := c.Request.RemoteAddr

	atomic.AddInt32(&state.Stats.Requests, 1)

	if status == StatusSealed {
		log.Printf("Rejected token request from %s. Reason: sealed.", remoteIp)
		atomic.AddInt32(&state.Stats.Denied, 1)
		c.JSON(503, struct {
			Status string `json:"status"`
			Ok     bool   `json:"ok"`
			Error  string `json:"error"`
		}{string(state.Status), false, "Gatekeeper is sealed."})
		return
	}

	var reqParams struct {
		TaskId string `json:"task_id"`
	}
	decoder := json.NewDecoder(c.Request.Body)
	if err := decoder.Decode(&reqParams); err == nil {
		if usedTaskIds.Has(reqParams.TaskId) {
			log.Printf("Rejected token request from %s (Task Id: %s). Reason: %v", remoteIp, reqParams.TaskId, errAlreadyGivenKey)
			atomic.AddInt32(&state.Stats.Denied, 1)
			c.JSON(403, struct {
				Status string `json:"status"`
				Ok     bool   `json:"ok"`
				Error  string `json:"error"`
			}{string(state.Status), false, errAlreadyGivenKey.Error()})
			return
		}
		if task, err := getMesosTask(reqParams.TaskId); err == nil {
			if len(task.Statuses) == 0 {
				log.Printf("Rejected token request from %s (Task Id: %s). Reason: %v (no status)", remoteIp, reqParams.TaskId, errTaskNotFresh)
				atomic.AddInt32(&state.Stats.Denied, 1)
				c.JSON(403, struct {
					Status string `json:"status"`
					Ok     bool   `json:"ok"`
					Error  string `json:"error"`
				}{string(state.Status), false, errTaskNotFresh.Error()})
				return
			}
			startTime := time.Unix(0, int64(task.Statuses[len(task.Statuses)-1].Timestamp*1000000000))
			if time.Now().Sub(startTime) > config.MaxTaskLife {
				log.Printf("Rejected token request from %s (Task Id: %s). Reason: %v (no status)", remoteIp, reqParams.TaskId, errTaskNotFresh)
				atomic.AddInt32(&state.Stats.Denied, 1)
				c.JSON(403, struct {
					Status string `json:"status"`
					Ok     bool   `json:"ok"`
					Error  string `json:"error"`
				}{string(state.Status), false, errTaskNotFresh.Error()})
				return
			}
			state.RLock()
			policy := activePolicies.Get(task.Name)
			state.RUnlock()
			if tempToken, err := createTokenPair(token, policy); err == nil {
				log.Printf("Provided token pair for %s (Task Id: %s) (Task Name: %s). Policies: %v", remoteIp, reqParams.TaskId, task.Name, policy.Policies)
				atomic.AddInt32(&state.Stats.Successful, 1)
				usedTaskIds.Put(reqParams.TaskId, config.MaxTaskLife+1*time.Minute)
				c.JSON(200, struct {
					Status string `json:"status"`
					Ok     bool   `json:"ok"`
					Token  string `json:"token"`
				}{string(state.Status), true, tempToken})
			} else {
				log.Printf("Failed to create token pair for %s (Task Id: %s). Reason: %v", remoteIp, reqParams.TaskId, errTaskNotFresh)
				atomic.AddInt32(&state.Stats.Denied, 1)
				c.JSON(500, struct {
					Status string `json:"status"`
					Ok     bool   `json:"ok"`
					Error  string `json:"error"`
				}{string(state.Status), false, err.Error()})
			}
		} else if err == errNoSuchTask {
			log.Printf("Rejected token request from %s (Task Id: %s). Reason: %v", remoteIp, reqParams.TaskId, errNoSuchTask)
			atomic.AddInt32(&state.Stats.Denied, 1)
			c.JSON(403, struct {
				Status string `json:"status"`
				Ok     bool   `json:"ok"`
				Error  string `json:"error"`
			}{string(state.Status), false, err.Error()})
		} else {
			log.Printf("Failed to retrieve task information for %s (Task Id: %s). Reason: %v", remoteIp, reqParams.TaskId, err)
			atomic.AddInt32(&state.Stats.Denied, 1)
			c.JSON(500, struct {
				Status string `json:"status"`
				Ok     bool   `json:"ok"`
				Error  string `json:"error"`
			}{string(state.Status), false, err.Error()})
		}
	} else {
		log.Printf("Rejected token request from %s. Reason: %v", remoteIp, err)
		atomic.AddInt32(&state.Stats.Denied, 1)
		c.JSON(400, struct {
			Status string `json:"status"`
			Ok     bool   `json:"ok"`
			Error  string `json:"error"`
		}{string(state.Status), false, err.Error()})
	}
}
<commit_msg>Retry getting the task information if the task currently has no status in mesos.<commit_after>package main

import (
	"encoding/json"
	"errors"
	"github.com/franela/goreq"
	"github.com/gin-gonic/gin"
	"log"
	"sync/atomic"
	"time"
)

var errTaskNotFresh = errors.New("This task has been running too long to request a token.")
var errAlreadyGivenKey = errors.New("This task has already been given a token.")
var usedTaskIds = NewTtlSet()

func createToken(token string, opts interface{}) (string, error) {
	r, err := goreq.Request{
		Uri:    vaultPath("/v1/auth/token/create", ""),
		Method: "POST",
		Body:   opts,
	}.WithHeader("X-Vault-Token", token).Do()
	if err == nil {
		defer r.Body.Close()
		switch r.StatusCode {
		case 200:
			var t vaultTokenResp
			if err := r.Body.FromJsonTo(&t); err == nil {
				return t.Auth.ClientToken, nil
			} else {
				return "", err
			}
		default:
			var e vaultError
			e.Code = r.StatusCode
			if err := r.Body.FromJsonTo(&e); err == nil {
				return "", e
			} else {
				e.Errors = []string{"communication error."}
				return "", e
			}
		}
	} else {
		return "", err
	}
}

func createTokenPair(token string, p *policy) (string, error) {
	tempTokenOpts := struct {
		Ttl      string   `json:"ttl"`
		NumUses  int      `json:"num_uses"`
		Policies []string `json:"policies"`
		NoParent bool     `json:"no_parent"`
	}{"10m", 2, []string{"default"}, true}
	pol := p.Policies
	if len(pol) == 0 { // explicitly set the policy, else the token will inherit ours
		pol = []string{"default"}
	}
	permTokenOpts := struct {
		Ttl      string            `json:"ttl,omitempty"`
		Policies []string          `json:"policies"`
		Meta     map[string]string `json:"meta,omitempty"`
		NumUses  int               `json:"num_uses"`
		NoParent bool              `json:"no_parent"`
	}{time.Duration(time.Duration(p.Ttl) * time.Second).String(), pol, p.Meta, p.NumUses, true}

	if tempToken, err := createToken(token, tempTokenOpts); err == nil {
		if permToken, err := createToken(token, permTokenOpts); err == nil {
			r, err := goreq.Request{
				Uri:    vaultPath("/v1/cubbyhole/vault-token", ""),
				Method: "POST",
				Body: struct {
					Token string `json:"token"`
				}{permToken},
			}.WithHeader("X-Vault-Token", tempToken).Do()
			if err == nil {
				defer r.Body.Close()
				switch r.StatusCode {
				case 204:
					return tempToken, nil
				default:
					var e vaultError
					e.Code = r.StatusCode
					if err := r.Body.FromJsonTo(&e); err == nil {
						return "", e
					} else {
						e.Errors = []string{"communication error."}
						return "", e
					}
				}
			} else {
				return "", err
			}
		} else {
			return "", err
		}
	} else {
		return "", err
	}
}

func Provide(c *gin.Context) {
	state.RLock()
	status := state.Status
	token := state.Token
	state.RUnlock()

	remoteIp := c.Request.RemoteAddr

	atomic.AddInt32(&state.Stats.Requests, 1)

	if status == StatusSealed {
		log.Printf("Rejected token request from %s. Reason: sealed.", remoteIp)
		atomic.AddInt32(&state.Stats.Denied, 1)
		c.JSON(503, struct {
			Status string `json:"status"`
			Ok     bool   `json:"ok"`
			Error  string `json:"error"`
		}{string(state.Status), false, "Gatekeeper is sealed."})
		return
	}

	var reqParams struct {
		TaskId string `json:"task_id"`
	}
	decoder := json.NewDecoder(c.Request.Body)
	if err := decoder.Decode(&reqParams); err == nil {
		if usedTaskIds.Has(reqParams.TaskId) {
			log.Printf("Rejected token request from %s (Task Id: %s). Reason: %v", remoteIp, reqParams.TaskId, errAlreadyGivenKey)
			atomic.AddInt32(&state.Stats.Denied, 1)
			c.JSON(403, struct {
				Status string `json:"status"`
				Ok     bool   `json:"ok"`
				Error  string `json:"error"`
			}{string(state.Status), false, errAlreadyGivenKey.Error()})
			return
		}
		/*
			Sometimes the mesos task status isn't available yet in mesos
			when we are asked for a token. In this case we wait a little while
			and then try to get the task info again.
		*/
		gMT := func(taskId string) (mesosTask, error) {
			task, err := getMesosTask(taskId)
			if err == nil && len(task.Statuses) == 0 {
				time.Sleep(500 * time.Millisecond)
				task, err = getMesosTask(taskId)
			}
			return task, err
		}
		if task, err := gMT(reqParams.TaskId); err == nil {
			if len(task.Statuses) == 0 {
				log.Printf("Rejected token request from %s (Task Id: %s). Reason: %v (no status)", remoteIp, reqParams.TaskId, errTaskNotFresh)
				atomic.AddInt32(&state.Stats.Denied, 1)
				c.JSON(403, struct {
					Status string `json:"status"`
					Ok     bool   `json:"ok"`
					Error  string `json:"error"`
				}{string(state.Status), false, errTaskNotFresh.Error()})
				return
			}
			startTime := time.Unix(0, int64(task.Statuses[len(task.Statuses)-1].Timestamp*1000000000))
			if time.Now().Sub(startTime) > config.MaxTaskLife {
				log.Printf("Rejected token request from %s (Task Id: %s). Reason: %v (no status)", remoteIp, reqParams.TaskId, errTaskNotFresh)
				atomic.AddInt32(&state.Stats.Denied, 1)
				c.JSON(403, struct {
					Status string `json:"status"`
					Ok     bool   `json:"ok"`
					Error  string `json:"error"`
				}{string(state.Status), false, errTaskNotFresh.Error()})
				return
			}
			state.RLock()
			policy := activePolicies.Get(task.Name)
			state.RUnlock()
			if tempToken, err := createTokenPair(token, policy); err == nil {
				log.Printf("Provided token pair for %s (Task Id: %s) (Task Name: %s). Policies: %v", remoteIp, reqParams.TaskId, task.Name, policy.Policies)
				atomic.AddInt32(&state.Stats.Successful, 1)
				usedTaskIds.Put(reqParams.TaskId, config.MaxTaskLife+1*time.Minute)
				c.JSON(200, struct {
					Status string `json:"status"`
					Ok     bool   `json:"ok"`
					Token  string `json:"token"`
				}{string(state.Status), true, tempToken})
			} else {
				log.Printf("Failed to create token pair for %s (Task Id: %s). Reason: %v", remoteIp, reqParams.TaskId, errTaskNotFresh)
				atomic.AddInt32(&state.Stats.Denied, 1)
				c.JSON(500, struct {
					Status string `json:"status"`
					Ok     bool   `json:"ok"`
					Error  string `json:"error"`
				}{string(state.Status), false, err.Error()})
			}
		} else if err == errNoSuchTask {
			log.Printf("Rejected token request from %s (Task Id: %s). Reason: %v", remoteIp, reqParams.TaskId, errNoSuchTask)
			atomic.AddInt32(&state.Stats.Denied, 1)
			c.JSON(403, struct {
				Status string `json:"status"`
				Ok     bool   `json:"ok"`
				Error  string `json:"error"`
			}{string(state.Status), false, err.Error()})
		} else {
			log.Printf("Failed to retrieve task information for %s (Task Id: %s). Reason: %v", remoteIp, reqParams.TaskId, err)
			atomic.AddInt32(&state.Stats.Denied, 1)
			c.JSON(500, struct {
				Status string `json:"status"`
				Ok     bool   `json:"ok"`
				Error  string `json:"error"`
			}{string(state.Status), false, err.Error()})
		}
	} else {
		log.Printf("Rejected token request from %s. Reason: %v", remoteIp, err)
		atomic.AddInt32(&state.Stats.Denied, 1)
		c.JSON(400, struct {
			Status string `json:"status"`
			Ok     bool   `json:"ok"`
			Error  string `json:"error"`
		}{string(state.Status), false, err.Error()})
	}
}
<|endoftext|>


<commit_before>package logberry

import (
	"sync/atomic"
)

// Task represents a particular component, function, or activity.  In
// general a Task is meant to be used within a single thread of
// execution, and the calling code is responsible for managing any
// concurrent manipulation.
type Task struct {
	uid uint64

	root *Root

	parent *Task

	component string

	activity string

	data D
}

var numtasks uint64

func newtaskuid() uint64 {
	// We have seen this atomic call cause problems on ARM...
	return atomic.AddUint64(&numtasks, 1) - 1
}

func newtask(parent *Task, component string, activity string, data []interface{}) *Task {

	t := &Task{
		uid:    newtaskuid(),
		parent: parent,
		activity: activity,
		data: DAggregate(data),
	}

	if parent != nil {
		t.root = parent.root
		t.component = parent.component
	} else {
		t.root = Std
	}

	if component != "" {
		t.component = component
	}

	t.root.event(t, BEGIN, t.activity+" begin", t.data)

	return t

}

// Task creates a new sub-task.  Parameter activity should be a short
// natural language description of the work that the Task represents,
// without any terminating punctuation.  Any data given here will be
// associated with the Task and reported with all its events.
func (x *Task) Task(activity string, data ...interface{}) *Task {
	return newtask(x, "", activity, data)
}

// Component creates a new Task object representing related long-lived
// functionality, rather than a directed, tightly scoped line of
// computation.  Parameter component should be a short lowercase
// string identifying the class, module, or other component that this
// Task represents.  The activity text of this Task is set to be
// "Component " + component.  Any data given will be associated with
// the Task and reported with all its events.
func (x *Task) Component(component string, data ...interface{}) *Task {
	return newtask(x, component, "Component " + component, data)
}

// AddData incorporates the given data into that associated and
// reported with this Task.  The rules for this construction are
// explained in CopyFrom.  This call does not generate a log event.
// The host Task is passed through as the return.  Among other things,
// this function is useful to silently accumulate data into the Task
// as it proceeds, to be reported when it concludes.
func (x *Task) AddData(data ...interface{}) *Task {
	x.data.CopyFrom(data)
	return x
}

// Event generates a user-specified log event.  Parameter event tags
// the class of the event, generally a short lowercase whitespace-free
// identifier.  A human-oriented text message is given as the msg
// parameter.  This should generally be static, short, use sentence
// capitalization but no terminating punctuation, and not itself
// include any data, which is better left to the structured data.  The
// variadic data parameter is aggregated as a D and reported with the
// event, as is the data permanently associated with the Task.  The
// given data is not associated to the Task permanently.
func (x *Task) Event(event string, msg string, data ...interface{}) {
	x.root.event(x, event, msg, DAggregate(data).CopyFrom(x.data))
}

// Info generates an informational log event.  A human-oriented text
// message is given as the msg parameter.  This should generally be
// static, short, use sentence capitalization but no terminating
// punctuation, and not itself include any data, which is better left
// to the structured data.  The variadic data parameter is aggregated
// as a D and reported with the event, as is the data permanently
// associated with the Task.  The given data is not associated to the
// Task permanently.
func (x *Task) Info(msg string, data ...interface{}) {
	x.root.event(x, INFO, msg, DAggregate(data).CopyFrom(x.data))
}

// Warning generates a warning log event indicating that a fault was
// encountered but the task is proceeding acceptably.  This should
// generally be static, short, use sentence capitalization but no
// terminating punctuation, and not itself include any data, which is
// better left to the structured data.  The variadic data parameter is
// aggregated as a D and reported with the event, as is the data
// permanently associated with the Task.  The given data is not
// associated to the Task permanently.
func (x *Task) Warning(msg string, data ...interface{}) {
	x.root.event(x, WARNING, msg, DAggregate(data).CopyFrom(x.data))
}

// Ready generates a ready log event reporting that the activity or
// component the Task represents is initialized and prepared to begin.
// The variadic data parameter is aggregated as a D and reported with
// the event, as is the data permanently associated with the Task.
// The given data is not associated to the Task permanently.
func (x *Task) Ready(data ...interface{}) {
	x.root.event(x, READY, x.activity+" ready",
		DAggregate(data).CopyFrom(x.data))
}

// Stopped generates a stopped log event reporting that the activity
// or component the Task represents has paused or shutdown.  The
// variadic data parameter is aggregated as a D and reported with the
// event, as is the data permanently associated with the Task.  The
// given data is not associated to the Task permanently.
func (x *Task) Stopped(data ...interface{}) {
	x.root.event(x, STOPPED, x.activity+" stopped",
		DAggregate(data).CopyFrom(x.data))
}

// Finalized generates an end log event reporting that the component
// the Task represents has ceased.  It is generally intended to be
// used for components, while Success is used for discrete activities.
// Continuing to use the Task is discouraged.  The variadic data
// parameter is aggregated as a D and reported with the event, as is
// the data permanently associated with the Task.  The given data is
// not associated to the Task permanently.
func (x *Task) Finalized(data ...interface{}) {
	x.root.event(x, END, x.activity+" finalized", DAggregate(data).CopyFrom(x.data))
}

// Success generates a success log event reporting that the activity
// the Task represents has concluded successfully.  It always returns
// nil.  Continuing to use the Task is discouraged.  The variadic data
// parameter is aggregated as a D and reported with the event, as is
// the data permanently associated with the Task.  The given data is
// not associated to the Task permanently.
func (x *Task) Success(data ...interface{}) error {
	x.root.event(x, SUCCESS, x.activity+" success", DAggregate(data).CopyFrom(x.data))
	return nil
}

// Error generates an error log event reporting an unrecoverable fault
// in an activity or component.  An error is returned wrapping the
// original error with a message reporting that the Task's activity
// has failed.  Continuing to use the Task is discouraged.  The
// variadic data parameter is aggregated as a D and embedded in the
// generated error.  It and the data permanently associated with the
// Task is reported with the event.  The reported source code position
// of the generated task error is adjusted to be the event invocation.
func (x *Task) Error(err error, data ...interface{}) error {
	m := x.activity + " failed"
	x.root.event(x, ERROR, m, D{"Error": DAggregate([]interface{}{err})}.CopyFrom(DAggregate(data)).CopyFrom(x.data))

	e := wraperror(m, err, data)
	e.Locate(1)
	return e
}

// Failure generates an error log event reporting an unrecoverable
// fault.  Failure and Error are essentially the same, the difference
// being that Failure is the first point of fault while Error takes an
// underlying error typically returned from another function or
// component.  An error is returned reporting that the activity or
// component represented by the Task has failed due to the underlying
// cause given in the message.  Continuing to use the Task is
// discouraged.  The variadic data parameter is aggregated as a D and
// embedded in the generated task error.  It and the data permanently
// associated with the Task is reported with the event.  The reported
// source code position of the generated task error is adjusted to be
// the event invocation.
func (x *Task) Failure(msg string, data ...interface{}) error {
	err := newerror(msg, nil)
	err.Locate(1)

	m := x.activity + " failed"
	x.root.event(x, ERROR, m, D{"Error": DAggregate([]interface{}{err})}.CopyFrom(DAggregate(data)).CopyFrom(x.data))

	e := wraperror(m, err, data)
	e.Locate(1)
	return e
}
<commit_msg>Realistically need WrapError().<commit_after>package logberry

import (
	"sync/atomic"
)

// Task represents a particular component, function, or activity.  In
// general a Task is meant to be used within a single thread of
// execution, and the calling code is responsible for managing any
// concurrent manipulation.
type Task struct {
	uid uint64

	root *Root

	parent *Task

	component string

	activity string

	data D
}

var numtasks uint64

func newtaskuid() uint64 {
	// We have seen this atomic call cause problems on ARM...
	return atomic.AddUint64(&numtasks, 1) - 1
}

func newtask(parent *Task, component string, activity string, data []interface{}) *Task {

	t := &Task{
		uid:    newtaskuid(),
		parent: parent,
		activity: activity,
		data: DAggregate(data),
	}

	if parent != nil {
		t.root = parent.root
		t.component = parent.component
	} else {
		t.root = Std
	}

	if component != "" {
		t.component = component
	}

	t.root.event(t, BEGIN, t.activity+" begin", t.data)

	return t

}

// Task creates a new sub-task.  Parameter activity should be a short
// natural language description of the work that the Task represents,
// without any terminating punctuation.  Any data given here will be
// associated with the Task and reported with all its events.
func (x *Task) Task(activity string, data ...interface{}) *Task {
	return newtask(x, "", activity, data)
}

// Component creates a new Task object representing related long-lived
// functionality, rather than a directed, tightly scoped line of
// computation.  Parameter component should be a short lowercase
// string identifying the class, module, or other component that this
// Task represents.  The activity text of this Task is set to be
// "Component " + component.  Any data given will be associated with
// the Task and reported with all its events.
func (x *Task) Component(component string, data ...interface{}) *Task {
	return newtask(x, component, "Component " + component, data)
}

// AddData incorporates the given data into that associated and
// reported with this Task.  The rules for this construction are
// explained in CopyFrom.  This call does not generate a log event.
// The host Task is passed through as the return.  Among other things,
// this function is useful to silently accumulate data into the Task
// as it proceeds, to be reported when it concludes.
func (x *Task) AddData(data ...interface{}) *Task {
	x.data.CopyFrom(data)
	return x
}

// Event generates a user-specified log event.  Parameter event tags
// the class of the event, generally a short lowercase whitespace-free
// identifier.  A human-oriented text message is given as the msg
// parameter.  This should generally be static, short, use sentence
// capitalization but no terminating punctuation, and not itself
// include any data, which is better left to the structured data.  The
// variadic data parameter is aggregated as a D and reported with the
// event, as is the data permanently associated with the Task.  The
// given data is not associated to the Task permanently.
func (x *Task) Event(event string, msg string, data ...interface{}) {
	x.root.event(x, event, msg, DAggregate(data).CopyFrom(x.data))
}

// Info generates an informational log event.  A human-oriented text
// message is given as the msg parameter.  This should generally be
// static, short, use sentence capitalization but no terminating
// punctuation, and not itself include any data, which is better left
// to the structured data.  The variadic data parameter is aggregated
// as a D and reported with the event, as is the data permanently
// associated with the Task.  The given data is not associated to the
// Task permanently.
func (x *Task) Info(msg string, data ...interface{}) {
	x.root.event(x, INFO, msg, DAggregate(data).CopyFrom(x.data))
}

// Warning generates a warning log event indicating that a fault was
// encountered but the task is proceeding acceptably.  This should
// generally be static, short, use sentence capitalization but no
// terminating punctuation, and not itself include any data, which is
// better left to the structured data.  The variadic data parameter is
// aggregated as a D and reported with the event, as is the data
// permanently associated with the Task.  The given data is not
// associated to the Task permanently.
func (x *Task) Warning(msg string, data ...interface{}) {
	x.root.event(x, WARNING, msg, DAggregate(data).CopyFrom(x.data))
}

// Ready generates a ready log event reporting that the activity or
// component the Task represents is initialized and prepared to begin.
// The variadic data parameter is aggregated as a D and reported with
// the event, as is the data permanently associated with the Task.
// The given data is not associated to the Task permanently.
func (x *Task) Ready(data ...interface{}) {
	x.root.event(x, READY, x.activity+" ready",
		DAggregate(data).CopyFrom(x.data))
}

// Stopped generates a stopped log event reporting that the activity
// or component the Task represents has paused or shutdown.  The
// variadic data parameter is aggregated as a D and reported with the
// event, as is the data permanently associated with the Task.  The
// given data is not associated to the Task permanently.
func (x *Task) Stopped(data ...interface{}) {
	x.root.event(x, STOPPED, x.activity+" stopped",
		DAggregate(data).CopyFrom(x.data))
}

// Finalized generates an end log event reporting that the component
// the Task represents has ceased.  It is generally intended to be
// used for components, while Success is used for discrete activities.
// Continuing to use the Task is discouraged.  The variadic data
// parameter is aggregated as a D and reported with the event, as is
// the data permanently associated with the Task.  The given data is
// not associated to the Task permanently.
func (x *Task) Finalized(data ...interface{}) {
	x.root.event(x, END, x.activity+" finalized", DAggregate(data).CopyFrom(x.data))
}

// Success generates a success log event reporting that the activity
// the Task represents has concluded successfully.  It always returns
// nil.  Continuing to use the Task is discouraged.  The variadic data
// parameter is aggregated as a D and reported with the event, as is
// the data permanently associated with the Task.  The given data is
// not associated to the Task permanently.
func (x *Task) Success(data ...interface{}) error {
	x.root.event(x, SUCCESS, x.activity+" success", DAggregate(data).CopyFrom(x.data))
	return nil
}

// Error generates an error log event reporting an unrecoverable fault
// in an activity or component.  An error is returned wrapping the
// original error with a message reporting that the Task's activity
// has failed.  Continuing to use the Task is discouraged.  The
// variadic data parameter is aggregated as a D and embedded in the
// generated error.  It and the data permanently associated with the
// Task is reported with the event.  The reported source code position
// of the generated task error is adjusted to be the event invocation.
func (x *Task) Error(err error, data ...interface{}) error {
	m := x.activity + " failed"
	x.root.event(x, ERROR, m, D{"Error": DAggregate([]interface{}{err})}.CopyFrom(DAggregate(data)).CopyFrom(x.data))

	e := wraperror(m, err, data)
	e.Locate(1)
	return e
}

func (x *Task) WrapError(msg string, source error, data ...interface{}) error {
	err := wraperror(msg, source, nil)
	err.Locate(1)

	m := x.activity + " failed"
	x.root.event(x, ERROR, m, D{"Error": DAggregate([]interface{}{err})}.CopyFrom(DAggregate(data)).CopyFrom(x.data))

	e := wraperror(m, err, data)
	e.Locate(1)
	return e
}

// Failure generates an error log event reporting an unrecoverable
// fault.  Failure and Error are essentially the same, the difference
// being that Failure is the first point of fault while Error takes an
// underlying error typically returned from another function or
// component.  An error is returned reporting that the activity or
// component represented by the Task has failed due to the underlying
// cause given in the message.  Continuing to use the Task is
// discouraged.  The variadic data parameter is aggregated as a D and
// embedded in the generated task error.  It and the data permanently
// associated with the Task is reported with the event.  The reported
// source code position of the generated task error is adjusted to be
// the event invocation.
func (x *Task) Failure(msg string, data ...interface{}) error {
	err := newerror(msg, nil)
	err.Locate(1)

	m := x.activity + " failed"
	x.root.event(x, ERROR, m, D{"Error": DAggregate([]interface{}{err})}.CopyFrom(DAggregate(data)).CopyFrom(x.data))

	e := wraperror(m, err, data)
	e.Locate(1)
	return e
}
<|endoftext|>


<commit_before>/* Copyright (c) 2014, Daniel Martí <mvdan@mvdan.cc> */
/* See LICENSE for licensing information */

package main

import (
	"crypto/rand"
	"encoding/hex"
	"errors"
	"flag"
	"fmt"
	"html/template"
	"io/ioutil"
	"log"
	"net/http"
	"os"
	"path"
	"path/filepath"
	"regexp"
	"strconv"
	"strings"
	"time"
)

const (
	indexTmpl = "index.html"
	formTmpl  = "form.html"
	idSize    = 8
	rawIdSize = idSize / 2
	randTries = 10
	fieldName = "paste"

	// GET error messages
	invalidId     = "Invalid paste id."
	pasteNotFound = "Paste doesn't exist."
	unknownError  = "Something went terribly wrong."

	// Common error messages
	timedOut = "Request timed out."
)

var (
	siteUrl, listen, dataDir, maxSizeStr string
	lifeTime, timeout                    time.Duration
	maxSize                              ByteSize
	indexTemplate, formTemplate          *template.Template

	regexByteSize = regexp.MustCompile(`^([\d\.]+)\s*([KM]?B|[BKM])$`)
	startTime     = time.Now()
)

var defMimeType = "text-plain; charset=utf-8"
var mimeTypes = map[string]string{
	".png":  "image/png",
	".gif":  "image/gif",
	".jpg":  "image/jpeg",
	".jpeg": "image/jpeg",
	".svg":  "image/svg+xml",
	".ico":  "image/x-icon",
	".pdf":  "application/pdf",
}

var workers [256]Worker
var post = make(chan PostRequest) // Posting is shared to balance load

type Id [rawIdSize]byte

type PasteInfo struct {
	Etag, ContentType, Path, Expires string
	ModTime                          time.Time
}

type GetRequest struct {
	w    http.ResponseWriter
	r    *http.Request
	done chan struct{}
	id   Id
}

type PostRequest struct {
	w       http.ResponseWriter
	r       *http.Request
	done    chan struct{}
	content []byte
	ext     string
	modTime time.Time
}

type Worker struct {
	n   byte // Its number, aka the first two hex chars
	get chan GetRequest
	del chan Id
	m   map[Id]PasteInfo
}

func (w Worker) recoverPaste(filePath string, fileInfo os.FileInfo, err error) error {
	if err != nil {
		return err
	}
	if fileInfo.IsDir() {
		return nil
	}
	dirParts := strings.Split(filePath, string(filepath.Separator))
	if len(dirParts) != 2 {
		return errors.New("Found invalid number of directories at " + filePath)
	}
	hexId := dirParts[0] + dirParts[1]
	ext := filepath.Ext(hexId)
	id, err := IdFromString(hexId[:len(hexId)-len(ext)])
	if err != nil {
		return err
	}
	modTime := fileInfo.ModTime()
	deathTime := modTime.Add(lifeTime)
	if deathTime.Before(startTime) {
		return os.Remove(filePath)
	}
	if modTime.After(startTime) {
		modTime = startTime
	}
	w.m[id] = id.GenPasteInfo(modTime, ext)
	w.DeletePasteAfter(id, deathTime.Sub(startTime))
	return nil
}

func (w Worker) RandomId() (id Id, err error) {
	id[0] = w.n
	for try := 0; try < randTries; try++ {
		if _, err := rand.Read(id[1:]); err != nil {
			return
		}
		if _, e := w.m[id]; !e {
			return id, nil
		}
	}
	return id, fmt.Errorf("Gave up trying to find an unused random id after %d tries", randTries)
}

func (w Worker) Work() {
	dir := hex.EncodeToString([]byte{w.n})
	if stat, err := os.Stat(dir); err == nil {
		if !stat.IsDir() {
			log.Fatalf("%s/%s exists but is not a directory!", dataDir, dir)
		}
	} else if err := os.Mkdir(dir, 0700); err != nil {
		log.Fatalf("Could not create data directory %s/%s: %s", dataDir, dir, err)
	}
	w.m = make(map[Id]PasteInfo)
	if err := filepath.Walk(dir, w.recoverPaste); err != nil {
		log.Fatalf("Could not recover data directory %s/%s: %s", dataDir, dir, err)
	}
	for {
		var done chan struct{}
		select {
		case request := <-w.get:
			done = request.done
			pasteInfo, e := w.m[request.id]
			if !e {
				http.Error(request.w, pasteNotFound, http.StatusNotFound)
				break
			}
			if inm := request.r.Header.Get("If-None-Match"); inm != "" {
				if pasteInfo.Etag == inm || inm == "*" {
					request.w.WriteHeader(http.StatusNotModified)
					break
				}
			}
			pasteFile, err := os.Open(pasteInfo.Path)
			if err != nil {
				http.Error(request.w, unknownError, http.StatusInternalServerError)
				break
			}
			request.w.Header().Set("Etag", pasteInfo.Etag)
			request.w.Header().Set("Expires", pasteInfo.Expires)
			request.w.Header().Set("Content-Type", pasteInfo.ContentType)
			http.ServeContent(request.w, request.r, "", pasteInfo.ModTime, pasteFile)
			pasteFile.Close()

		case request := <-post:
			done = request.done
			id, err := w.RandomId()
			if err != nil {
				log.Println(err)
				http.Error(request.w, unknownError, http.StatusInternalServerError)
				break
			}
			pasteInfo := id.GenPasteInfo(request.modTime, request.ext)
			pasteFile, err := os.OpenFile(pasteInfo.Path, os.O_WRONLY|os.O_CREATE|os.O_EXCL, 0600)
			if err != nil {
				log.Printf("Could not create new paste file %s: %s", pasteInfo.Path, err)
				http.Error(request.w, unknownError, http.StatusInternalServerError)
				break
			}
			_, err = pasteFile.Write(request.content)
			pasteFile.Close()
			if err != nil {
				log.Printf("Could not write data into %s: %s", pasteInfo.Path, err)
				http.Error(request.w, unknownError, http.StatusInternalServerError)
				break
			}
			w.m[id] = pasteInfo
			w.DeletePasteAfter(id, lifeTime)
			fmt.Fprintf(request.w, "%s/%s\n", siteUrl, id)

		case id := <-w.del:
			pasteInfo, _ := w.m[id]
			if err := os.Remove(pasteInfo.Path); err == nil {
				delete(w.m, id)
			} else {
				log.Printf("Could not remove %s: %s", id, err)
				w.DeletePasteAfter(id, 2*time.Minute)
			}
		}
		if done != nil {
			done <- struct{}{}
		}
	}
}

func init() {
	flag.StringVar(&siteUrl, "u", "http://localhost:8080", "URL of the site")
	flag.StringVar(&listen, "l", "localhost:8080", "Host and port to listen to")
	flag.StringVar(&dataDir, "d", "data", "Directory to store all the pastes in")
	flag.DurationVar(&lifeTime, "t", 12*time.Hour, "Lifetime of the pastes")
	flag.DurationVar(&timeout, "T", 200*time.Millisecond, "Timeout of requests")
	flag.StringVar(&maxSizeStr, "s", "1M", "Maximum size of POSTs in bytes")
}

func IdFromString(hexId string) (Id, error) {
	var id Id
	if len(hexId) != idSize {
		return id, errors.New("Invalid id at " + hexId)
	}
	b, err := hex.DecodeString(hexId)
	if err != nil || len(b) != rawIdSize {
		return id, errors.New("Invalid id at " + hexId)
	}
	copy(id[:], b)
	return id, nil
}

func (id Id) String() string {
	return hex.EncodeToString(id[:])
}

func (id Id) GenPasteInfo(modTime time.Time, ext string) (pasteInfo PasteInfo) {
	pasteInfo.ModTime = modTime
	pasteInfo.Expires = modTime.Add(lifeTime).UTC().Format(http.TimeFormat)
	pasteInfo.Etag = fmt.Sprintf("%d-%s", pasteInfo.ModTime.Unix(), id)
	if pasteInfo.ContentType = mimeTypes[ext]; pasteInfo.ContentType == "" {
		pasteInfo.ContentType = defMimeType
	}
	hexId := id.String()
	pasteInfo.Path = path.Join(hexId[0:2], hexId[2:]+ext)
	return
}

func (w Worker) DeletePasteAfter(id Id, duration time.Duration) {
	timer := time.NewTimer(duration)
	go func() {
		<-timer.C
		w.del <- id
	}()
}

type ByteSize int64

const (
	B ByteSize = 1 << (10 * iota)
	KB
	MB
)

func parseByteSize(str string) (ByteSize, error) {
	if !regexByteSize.MatchString(str) {
		return 0, errors.New("Could not parse size in bytes")
	}
	parts := regexByteSize.FindStringSubmatch(str)
	size, _ := strconv.ParseFloat(string(parts[1]), 64)
	switch string(parts[2]) {
	case "KB", "K":
		size *= float64(KB)
	case "MB", "M":
		size *= float64(MB)
	}
	return ByteSize(size), nil
}

func (b ByteSize) String() string {
	switch {
	case b >= MB:
		return fmt.Sprintf("%.2fMB", float64(b)/float64(MB))
	case b >= KB:
		return fmt.Sprintf("%.2fKB", float64(b)/float64(KB))
	}
	return fmt.Sprintf("%dB", b)
}

func handler(w http.ResponseWriter, r *http.Request) {
	done := make(chan struct{})
	timer := time.NewTimer(timeout)
	switch r.Method {
	case "GET":
		switch r.URL.Path {
		case "/":
			indexTemplate.Execute(w, struct{ SiteUrl, LifeTime, FieldName string }{
				siteUrl, lifeTime.String(), fieldName})
			return
		case "/form":
			formTemplate.Execute(w, struct{ SiteUrl, LifeTime, FieldName string }{
				siteUrl, lifeTime.String(), fieldName})
			return
		}
		id, err := IdFromString(r.URL.Path[1:])
		if err != nil {
			http.Error(w, invalidId, http.StatusBadRequest)
			return
		}
		select {
		case <-timer.C:
			http.Error(w, timedOut, http.StatusRequestTimeout)
		case workers[id[0]].get <- GetRequest{id: id, w: w, r: r, done: done}:
			timer.Stop()
		}
	case "POST":
		r.Body = http.MaxBytesReader(w, r.Body, int64(maxSize))
		var content []byte
		ext := ""
		if value := r.FormValue("paste"); value != "" {
			content = []byte(value)
		} else if f, h, err := r.FormFile("paste"); err == nil {
			content, err = ioutil.ReadAll(f)
			ext = strings.ToLower(filepath.Ext(h.Filename))
			f.Close()
			if err != nil {
				http.Error(w, err.Error(), http.StatusBadRequest)
				return
			}
		} else {
			http.Error(w, err.Error(), http.StatusBadRequest)
			return
		}
		select {
		case <-timer.C:
			http.Error(w, timedOut, http.StatusRequestTimeout)
		case post <- PostRequest{content: content, ext: ext, modTime: time.Now(), w: w, r: r, done: done}:
			timer.Stop()
		}
	default:
		http.Error(w, "Unsupported action.", http.StatusBadRequest)
		return
	}
	<-done
}

func main() {
	var err error
	flag.Parse()
	if maxSize, err = parseByteSize(maxSizeStr); err != nil {
		log.Fatalf("Invalid max size '%s': %s", maxSizeStr, err)
	}
	if indexTemplate, err = template.ParseFiles(indexTmpl); err != nil {
		log.Fatalf("Could not load template %s: %s", indexTmpl, err)
	}
	if formTemplate, err = template.ParseFiles(formTmpl); err != nil {
		log.Fatalf("Could not load template %s: %s", formTmpl, err)
	}
	if err = os.MkdirAll(dataDir, 0700); err != nil {
		log.Fatalf("Could not create data directory %s: %s", dataDir, err)
	}
	if err = os.Chdir(dataDir); err != nil {
		log.Fatalf("Could not enter data directory %s: %s", dataDir, err)
	}
	log.Printf("maxSize  = %s", maxSize)
	log.Printf("siteUrl  = %s", siteUrl)
	log.Printf("listen   = %s", listen)
	log.Printf("dataDir  = %s", dataDir)
	log.Printf("lifeTime = %s", lifeTime)
	log.Printf("timeout  = %s", timeout)
	for n := range workers {
		w := &workers[n]
		w.n = byte(n)
		w.get = make(chan GetRequest)
		w.del = make(chan Id)
		go w.Work()
	}
	http.HandleFunc("/", handler)
	log.Printf("Up and running!")
	log.Fatal(http.ListenAndServe(listen, nil))
}
<commit_msg>Actually use fieldName in the server code<commit_after>/* Copyright (c) 2014, Daniel Martí <mvdan@mvdan.cc> */
/* See LICENSE for licensing information */

package main

import (
	"crypto/rand"
	"encoding/hex"
	"errors"
	"flag"
	"fmt"
	"html/template"
	"io/ioutil"
	"log"
	"net/http"
	"os"
	"path"
	"path/filepath"
	"regexp"
	"strconv"
	"strings"
	"time"
)

const (
	indexTmpl = "index.html"
	formTmpl  = "form.html"
	idSize    = 8
	rawIdSize = idSize / 2
	randTries = 10
	fieldName = "paste"

	// GET error messages
	invalidId     = "Invalid paste id."
	pasteNotFound = "Paste doesn't exist."
	unknownError  = "Something went terribly wrong."

	// Common error messages
	timedOut = "Request timed out."
)

var (
	siteUrl, listen, dataDir, maxSizeStr string
	lifeTime, timeout                    time.Duration
	maxSize                              ByteSize
	indexTemplate, formTemplate          *template.Template

	regexByteSize = regexp.MustCompile(`^([\d\.]+)\s*([KM]?B|[BKM])$`)
	startTime     = time.Now()
)

var defMimeType = "text-plain; charset=utf-8"
var mimeTypes = map[string]string{
	".png":  "image/png",
	".gif":  "image/gif",
	".jpg":  "image/jpeg",
	".jpeg": "image/jpeg",
	".svg":  "image/svg+xml",
	".ico":  "image/x-icon",
	".pdf":  "application/pdf",
}

var workers [256]Worker
var post = make(chan PostRequest) // Posting is shared to balance load

type Id [rawIdSize]byte

type PasteInfo struct {
	Etag, ContentType, Path, Expires string
	ModTime                          time.Time
}

type GetRequest struct {
	w    http.ResponseWriter
	r    *http.Request
	done chan struct{}
	id   Id
}

type PostRequest struct {
	w       http.ResponseWriter
	r       *http.Request
	done    chan struct{}
	content []byte
	ext     string
	modTime time.Time
}

type Worker struct {
	n   byte // Its number, aka the first two hex chars
	get chan GetRequest
	del chan Id
	m   map[Id]PasteInfo
}

func (w Worker) recoverPaste(filePath string, fileInfo os.FileInfo, err error) error {
	if err != nil {
		return err
	}
	if fileInfo.IsDir() {
		return nil
	}
	dirParts := strings.Split(filePath, string(filepath.Separator))
	if len(dirParts) != 2 {
		return errors.New("Found invalid number of directories at " + filePath)
	}
	hexId := dirParts[0] + dirParts[1]
	ext := filepath.Ext(hexId)
	id, err := IdFromString(hexId[:len(hexId)-len(ext)])
	if err != nil {
		return err
	}
	modTime := fileInfo.ModTime()
	deathTime := modTime.Add(lifeTime)
	if deathTime.Before(startTime) {
		return os.Remove(filePath)
	}
	if modTime.After(startTime) {
		modTime = startTime
	}
	w.m[id] = id.GenPasteInfo(modTime, ext)
	w.DeletePasteAfter(id, deathTime.Sub(startTime))
	return nil
}

func (w Worker) RandomId() (id Id, err error) {
	id[0] = w.n
	for try := 0; try < randTries; try++ {
		if _, err := rand.Read(id[1:]); err != nil {
			return id, err
		}
		if _, e := w.m[id]; !e {
			return id, nil
		}
	}
	return id, fmt.Errorf("Gave up trying to find an unused random id after %d tries", randTries)
}

func (w Worker) Work() {
	dir := hex.EncodeToString([]byte{w.n})
	if stat, err := os.Stat(dir); err == nil {
		if !stat.IsDir() {
			log.Fatalf("%s/%s exists but is not a directory!", dataDir, dir)
		}
	} else if err := os.Mkdir(dir, 0700); err != nil {
		log.Fatalf("Could not create data directory %s/%s: %s", dataDir, dir, err)
	}
	w.m = make(map[Id]PasteInfo)
	if err := filepath.Walk(dir, w.recoverPaste); err != nil {
		log.Fatalf("Could not recover data directory %s/%s: %s", dataDir, dir, err)
	}
	for {
		var done chan struct{}
		select {
		case request := <-w.get:
			done = request.done
			pasteInfo, e := w.m[request.id]
			if !e {
				http.Error(request.w, pasteNotFound, http.StatusNotFound)
				break
			}
			if inm := request.r.Header.Get("If-None-Match"); inm != "" {
				if pasteInfo.Etag == inm || inm == "*" {
					request.w.WriteHeader(http.StatusNotModified)
					break
				}
			}
			pasteFile, err := os.Open(pasteInfo.Path)
			if err != nil {
				http.Error(request.w, unknownError, http.StatusInternalServerError)
				break
			}
			request.w.Header().Set("Etag", pasteInfo.Etag)
			request.w.Header().Set("Expires", pasteInfo.Expires)
			request.w.Header().Set("Content-Type", pasteInfo.ContentType)
			http.ServeContent(request.w, request.r, "", pasteInfo.ModTime, pasteFile)
			pasteFile.Close()

		case request := <-post:
			done = request.done
			id, err := w.RandomId()
			if err != nil {
				log.Println(err)
				http.Error(request.w, unknownError, http.StatusInternalServerError)
				break
			}
			pasteInfo := id.GenPasteInfo(request.modTime, request.ext)
			pasteFile, err := os.OpenFile(pasteInfo.Path, os.O_WRONLY|os.O_CREATE|os.O_EXCL, 0600)
			if err != nil {
				log.Printf("Could not create new paste file %s: %s", pasteInfo.Path, err)
				http.Error(request.w, unknownError, http.StatusInternalServerError)
				break
			}
			_, err = pasteFile.Write(request.content)
			pasteFile.Close()
			if err != nil {
				log.Printf("Could not write data into %s: %s", pasteInfo.Path, err)
				http.Error(request.w, unknownError, http.StatusInternalServerError)
				break
			}
			w.m[id] = pasteInfo
			w.DeletePasteAfter(id, lifeTime)
			fmt.Fprintf(request.w, "%s/%s\n", siteUrl, id)

		case id := <-w.del:
			pasteInfo, _ := w.m[id]
			if err := os.Remove(pasteInfo.Path); err == nil {
				delete(w.m, id)
			} else {
				log.Printf("Could not remove %s: %s", id, err)
				w.DeletePasteAfter(id, 2*time.Minute)
			}
		}
		if done != nil {
			done <- struct{}{}
		}
	}
}

func init() {
	flag.StringVar(&siteUrl, "u", "http://localhost:8080", "URL of the site")
	flag.StringVar(&listen, "l", "localhost:8080", "Host and port to listen to")
	flag.StringVar(&dataDir, "d", "data", "Directory to store all the pastes in")
	flag.DurationVar(&lifeTime, "t", 12*time.Hour, "Lifetime of the pastes")
	flag.DurationVar(&timeout, "T", 200*time.Millisecond, "Timeout of requests")
	flag.StringVar(&maxSizeStr, "s", "1M", "Maximum size of POSTs in bytes")
}

func IdFromString(hexId string) (Id, error) {
	var id Id
	if len(hexId) != idSize {
		return id, errors.New("Invalid id at " + hexId)
	}
	b, err := hex.DecodeString(hexId)
	if err != nil || len(b) != rawIdSize {
		return id, errors.New("Invalid id at " + hexId)
	}
	copy(id[:], b)
	return id, nil
}

func (id Id) String() string {
	return hex.EncodeToString(id[:])
}

func (id Id) GenPasteInfo(modTime time.Time, ext string) (pasteInfo PasteInfo) {
	pasteInfo.ModTime = modTime
	pasteInfo.Expires = modTime.Add(lifeTime).UTC().Format(http.TimeFormat)
	pasteInfo.Etag = fmt.Sprintf("%d-%s", pasteInfo.ModTime.Unix(), id)
	if pasteInfo.ContentType = mimeTypes[ext]; pasteInfo.ContentType == "" {
		pasteInfo.ContentType = defMimeType
	}
	hexId := id.String()
	pasteInfo.Path = path.Join(hexId[0:2], hexId[2:]+ext)
	return
}

func (w Worker) DeletePasteAfter(id Id, duration time.Duration) {
	timer := time.NewTimer(duration)
	go func() {
		<-timer.C
		w.del <- id
	}()
}

type ByteSize int64

const (
	B ByteSize = 1 << (10 * iota)
	KB
	MB
)

func parseByteSize(str string) (ByteSize, error) {
	if !regexByteSize.MatchString(str) {
		return 0, errors.New("Could not parse size in bytes")
	}
	parts := regexByteSize.FindStringSubmatch(str)
	size, _ := strconv.ParseFloat(string(parts[1]), 64)
	switch string(parts[2]) {
	case "KB", "K":
		size *= float64(KB)
	case "MB", "M":
		size *= float64(MB)
	}
	return ByteSize(size), nil
}

func (b ByteSize) String() string {
	switch {
	case b >= MB:
		return fmt.Sprintf("%.2fMB", float64(b)/float64(MB))
	case b >= KB:
		return fmt.Sprintf("%.2fKB", float64(b)/float64(KB))
	}
	return fmt.Sprintf("%dB", b)
}

func handler(w http.ResponseWriter, r *http.Request) {
	done := make(chan struct{})
	timer := time.NewTimer(timeout)
	switch r.Method {
	case "GET":
		switch r.URL.Path {
		case "/":
			indexTemplate.Execute(w, struct{ SiteUrl, LifeTime, FieldName string }{
				siteUrl, lifeTime.String(), fieldName})
			return
		case "/form":
			formTemplate.Execute(w, struct{ SiteUrl, LifeTime, FieldName string }{
				siteUrl, lifeTime.String(), fieldName})
			return
		}
		id, err := IdFromString(r.URL.Path[1:])
		if err != nil {
			http.Error(w, invalidId, http.StatusBadRequest)
			return
		}
		select {
		case <-timer.C:
			http.Error(w, timedOut, http.StatusRequestTimeout)
		case workers[id[0]].get <- GetRequest{id: id, w: w, r: r, done: done}:
			timer.Stop()
		}
	case "POST":
		r.Body = http.MaxBytesReader(w, r.Body, int64(maxSize))
		var content []byte
		ext := ""
		if value := r.FormValue(fieldName); value != "" {
			content = []byte(value)
		} else if f, h, err := r.FormFile(fieldName); err == nil {
			content, err = ioutil.ReadAll(f)
			ext = strings.ToLower(filepath.Ext(h.Filename))
			f.Close()
			if err != nil {
				http.Error(w, err.Error(), http.StatusBadRequest)
				return
			}
		} else {
			http.Error(w, err.Error(), http.StatusBadRequest)
			return
		}
		select {
		case <-timer.C:
			http.Error(w, timedOut, http.StatusRequestTimeout)
		case post <- PostRequest{content: content, ext: ext, modTime: time.Now(), w: w, r: r, done: done}:
			timer.Stop()
		}
	default:
		http.Error(w, "Unsupported action.", http.StatusBadRequest)
		return
	}
	<-done
}

func main() {
	var err error
	flag.Parse()
	if maxSize, err = parseByteSize(maxSizeStr); err != nil {
		log.Fatalf("Invalid max size '%s': %s", maxSizeStr, err)
	}
	if indexTemplate, err = template.ParseFiles(indexTmpl); err != nil {
		log.Fatalf("Could not load template %s: %s", indexTmpl, err)
	}
	if formTemplate, err = template.ParseFiles(formTmpl); err != nil {
		log.Fatalf("Could not load template %s: %s", formTmpl, err)
	}
	if err = os.MkdirAll(dataDir, 0700); err != nil {
		log.Fatalf("Could not create data directory %s: %s", dataDir, err)
	}
	if err = os.Chdir(dataDir); err != nil {
		log.Fatalf("Could not enter data directory %s: %s", dataDir, err)
	}
	log.Printf("maxSize  = %s", maxSize)
	log.Printf("siteUrl  = %s", siteUrl)
	log.Printf("listen   = %s", listen)
	log.Printf("dataDir  = %s", dataDir)
	log.Printf("lifeTime = %s", lifeTime)
	log.Printf("timeout  = %s", timeout)
	for n := range workers {
		w := &workers[n]
		w.n = byte(n)
		w.get = make(chan GetRequest)
		w.del = make(chan Id)
		go w.Work()
	}
	http.HandleFunc("/", handler)
	log.Printf("Up and running!")
	log.Fatal(http.ListenAndServe(listen, nil))
}
<|endoftext|>


<commit_before>/* Copyright (c) 2014, Daniel Martí <mvdan@mvdan.cc> */
/* See LICENSE for licensing information */

package main

import (
	"crypto/rand"
	"encoding/hex"
	"errors"
	"flag"
	"fmt"
	"html/template"
	"io"
	"log"
	"net/http"
	"os"
	"path"
	"path/filepath"
	"regexp"
	"strconv"
	"strings"
	"sync"
	"time"
)

const (
	indexTmpl = "index.html"
	formTmpl  = "form.html"
	idSize    = 8
	randTries = 10

	// GET error messages
	invalidId     = "Invalid paste id."
	pasteNotFound = "Paste doesn't exist."
	unknownError  = "Something went terribly wrong."
	// POST error messages
	missingForm = "Paste could not be found inside the posted form."
)

var (
	siteUrl, listen, dataDir string
	lifeTime                 time.Duration
	maxSizeStr               string
	maxSize                  ByteSize

	regexByteSize = regexp.MustCompile(`^([\d\.]+)\s*([KM]?B|[BKM])$`)
	indexTemplate *template.Template
	formTemplate  *template.Template
	data          = struct {
		sync.RWMutex
		m map[Id]PasteInfo
	}{m: make(map[Id]PasteInfo)}
)

func init() {
	flag.StringVar(&siteUrl, "u", "http://localhost:8080", "URL of the site")
	flag.StringVar(&listen, "l", "localhost:8080", "Host and port to listen to")
	flag.StringVar(&dataDir, "d", "data", "Directory to store all the pastes in")
	flag.DurationVar(&lifeTime, "t", 12*time.Hour, "Lifetime of the pastes (units: s,m,h)")
	flag.StringVar(&maxSizeStr, "s", "1M", "Maximum size of POSTs in bytes (units: B,K,M)")
}

type PasteInfo struct {
	ModTime     time.Time
	DeathTime   time.Time
	Size        ByteSize
	Etag        string
	ContentType string
}

type Id [idSize / 2]byte

func IdFromString(hexId string) (Id, error) {
	var id Id
	if len(hexId) != idSize {
		return id, errors.New("Invalid id")
	}
	b, err := hex.DecodeString(hexId)
	if err != nil || len(b) != idSize/2 {
		return id, errors.New("Invalid id")
	}
	copy(id[:], b)
	return id, nil
}

func IdFromPath(idPath string) (Id, error) {
	var id Id
	parts := strings.Split(idPath, string(filepath.Separator))
	if len(parts) != 2 {
		return id, errors.New("Found invalid number of directories at " + idPath)
	}
	return IdFromString(parts[0] + parts[1])
}

func RandomId() (Id, error) {
	var id Id
	data.RLock()
	defer data.RUnlock()
	for try := 0; try < randTries; try++ {
		_, err := rand.Read(id[:])
		if err != nil {
			return id, err
		}
		if _, e := data.m[id]; !e {
			return id, nil
		}
	}
	return id, fmt.Errorf("Gave up trying to find an unused random id after %d tries", randTries)
}

func (id Id) String() string {
	return hex.EncodeToString(id[:])
}

func (id Id) Path() string {
	hexId := id.String()
	return path.Join(hexId[0:2], hexId[2:])
}

func (id Id) EndLife() {
	data.Lock()
	defer data.Unlock()
	err := os.Remove(id.Path())
	if err == nil {
		delete(data.m, id)
		log.Printf("Removed paste: %s", id)
	} else {
		log.Printf("Could not end the life of %s: %s", id, err)
		id.EndLifeAfter(2 * time.Minute)
	}
}

func (id Id) EndLifeAfter(duration time.Duration) {
	timer := time.NewTimer(duration)
	go func() {
		<-timer.C
		id.EndLife()
	}()
}

type ByteSize int64

const (
	B ByteSize = 1 << (10 * iota)
	KB
	MB
)

func parseByteSize(str string) (ByteSize, error) {
	if !regexByteSize.MatchString(str) {
		return 0, errors.New("Could not parse size in bytes")
	}
	parts := regexByteSize.FindStringSubmatch(str)
	size, _ := strconv.ParseFloat(string(parts[1]), 64)

	switch string(parts[2]) {
	case "KB", "K":
		size *= float64(KB)
	case "MB", "M":
		size *= float64(MB)
	}
	return ByteSize(size), nil
}

func (b ByteSize) String() string {
	switch {
	case b >= MB:
		return fmt.Sprintf("%.2f MB", float64(b)/float64(MB))
	case b >= KB:
		return fmt.Sprintf("%.2f KB", float64(b)/float64(KB))
	}
	return fmt.Sprintf("%d B", b)
}

func handler(w http.ResponseWriter, r *http.Request) {
	var err error
	switch r.Method {
	case "GET":
		switch r.URL.Path {
		case "/":
			indexTemplate.Execute(w, struct {
				SiteUrl, LifeTime string
			}{
				siteUrl,
				fmt.Sprintf("%g hours and %g minutes", lifeTime.Hours(),
					lifeTime.Minutes()-lifeTime.Hours()*60),
			})
			return
		case "/form":
			formTemplate.Execute(w, struct{ SiteUrl string }{siteUrl})
			return
		}
		id, err := IdFromString(r.URL.Path[1:])
		if err != nil {
			w.WriteHeader(http.StatusBadRequest)
			fmt.Fprintf(w, "%s\n", invalidId)
			return
		}
		data.RLock()
		defer data.RUnlock()
		pasteInfo, e := data.m[id]
		if !e {
			w.WriteHeader(http.StatusNotFound)
			fmt.Fprintf(w, "%s\n", pasteNotFound)
			return
		}
		if inm := r.Header.Get("If-None-Match"); inm != "" {
			if pasteInfo.Etag == inm || inm == "*" {
				w.WriteHeader(http.StatusNotModified)
				return
			}
		}
		pastePath := id.Path()
		pasteFile, err := os.Open(pastePath)
		if err != nil {
			w.WriteHeader(http.StatusInternalServerError)
			fmt.Fprintf(w, "%s\n", unknownError)
			return
		}
		defer pasteFile.Close()
		w.Header().Set("Etag", pasteInfo.Etag)
		w.Header().Set("Content-Type", pasteInfo.ContentType)
		http.ServeContent(w, r, "", pasteInfo.ModTime, pasteFile)

	case "POST":
		r.Body = http.MaxBytesReader(w, r.Body, int64(maxSize))
		var content string
		if err = r.ParseMultipartForm(int64(maxSize)); err != nil {
			log.Printf("Could not parse POST multipart form: %s", err)
			w.WriteHeader(http.StatusBadRequest)
			fmt.Fprintf(w, "%s\n", err)
			return
		}
		if vs, found := r.Form["paste"]; found && len(vs[0]) > 0 {
			content = vs[0]
		} else {
			w.WriteHeader(http.StatusBadRequest)
			fmt.Fprintf(w, "%s\n", missingForm)
			return
		}
		id, err := RandomId()
		if err == nil {
			data.Lock()
			defer data.Unlock()
		} else {
			log.Println(err)
			w.WriteHeader(http.StatusInternalServerError)
			fmt.Fprintf(w, "%s\n", unknownError)
			return
		}
		pastePath := id.Path()
		dir, _ := path.Split(pastePath)
		if err = os.Mkdir(dir, 0700); err != nil {
			log.Printf("Could not create directories leading to %s: %s", pastePath, err)
			w.WriteHeader(http.StatusInternalServerError)
			fmt.Fprintf(w, "%s\n", unknownError)
			return
		}
		deathTime := time.Now().Add(lifeTime)
		id.EndLifeAfter(lifeTime)
		pasteFile, err := os.OpenFile(pastePath, os.O_WRONLY|os.O_CREATE|os.O_EXCL, 0600)
		if err != nil {
			log.Printf("Could not create new paste file %s: %s", pastePath, err)
			w.WriteHeader(http.StatusInternalServerError)
			fmt.Fprintf(w, "%s\n", unknownError)
			return
		}
		defer pasteFile.Close()
		b, err := io.WriteString(pasteFile, content)
		if err != nil {
			log.Printf("Could not write data into %s: %s", pastePath, err)
			w.WriteHeader(http.StatusInternalServerError)
			fmt.Fprintf(w, "%s\n", unknownError)
			return
		}
		writtenSize := ByteSize(b)
		pasteInfo := PasteInfo{
			ModTime:   time.Now(),
			DeathTime: deathTime,
			Size:      writtenSize,
		}
		pasteInfo.Etag = fmt.Sprintf("%d-%s", pasteInfo.ModTime.Unix(), id)
		pasteInfo.ContentType = http.DetectContentType([]byte(content))
		if pasteInfo.ContentType == "application/octet-stream" {
			pasteInfo.ContentType = "text-plain; charset=utf-8"
		}
		data.m[id] = pasteInfo
		log.Printf("Created new paste %s (%s %s) to die at %s",
			id, pasteInfo.ContentType, pasteInfo.Size, pasteInfo.DeathTime)
		fmt.Fprintf(w, "%s/%s\n", siteUrl, id)
	}
}

func walkFunc(filePath string, fileInfo os.FileInfo, err error) error {
	if err != nil {
		return err
	}
	if fileInfo.IsDir() {
		return nil
	}
	id, err := IdFromPath(filePath)
	if err != nil {
		return errors.New("Found incompatible id at path " + filePath)
	}
	modTime := fileInfo.ModTime()
	deathTime := modTime.Add(lifeTime)
	now := time.Now()
	if deathTime.Before(now) {
		go id.EndLife()
		return nil
	}
	var lifeLeft time.Duration
	if deathTime.After(now.Add(lifeTime)) {
		lifeLeft = lifeTime
	} else {
		lifeLeft = deathTime.Sub(now)
	}
	size := ByteSize(fileInfo.Size())
	pasteFile, err := os.Open(filePath)
	if err != nil {
		return err
	}
	defer pasteFile.Close()
	read := make([]byte, 512)
	_, err = pasteFile.Read(read)
	if err != nil {
		return err
	}
	pasteInfo := PasteInfo{
		ModTime:   modTime,
		DeathTime: deathTime,
		Size:      size,
		Etag:      fmt.Sprintf("%d-%s", modTime.Unix(), id),
	}
	pasteInfo.ContentType = http.DetectContentType(read)
	if pasteInfo.ContentType == "application/octet-stream" {
		pasteInfo.ContentType = "text-plain; charset=utf-8"
	}
	data.m[id] = pasteInfo
	log.Printf("Recovered paste %s (%s %s) from %s has %s left",
		id, pasteInfo.ContentType, pasteInfo.Size, pasteInfo.ModTime, lifeLeft)
	id.EndLifeAfter(lifeLeft)
	return nil
}

func main() {
	var err error
	flag.Parse()
	if maxSize, err = parseByteSize(maxSizeStr); err != nil {
		log.Fatalf("Invalid max size '%s': %s", maxSizeStr, err)
	}
	if indexTemplate, err = template.ParseFiles(indexTmpl); err != nil {
		log.Fatalf("Could not load template %s: %s", indexTmpl, err)
	}
	if formTemplate, err = template.ParseFiles(formTmpl); err != nil {
		log.Fatalf("Could not load template %s: %s", formTmpl, err)
	}
	if err = os.MkdirAll(dataDir, 0700); err != nil {
		log.Fatalf("Could not create data directory %s: %s", dataDir, err)
	}
	if err = os.Chdir(dataDir); err != nil {
		log.Fatalf("Could not enter data directory %s: %s", dataDir, err)
	}
	log.Printf("maxSize  = %s", maxSize)
	log.Printf("siteUrl  = %s", siteUrl)
	log.Printf("listen   = %s", listen)
	log.Printf("dataDir  = %s", dataDir)
	log.Printf("lifeTime = %s", lifeTime)
	if err = filepath.Walk(".", walkFunc); err != nil {
		log.Fatalf("Could not recover data directory %s: %s", dataDir, err)
	}
	http.HandleFunc("/", handler)
	log.Printf("Up and running!")
	log.Fatal(http.ListenAndServe(listen, nil))
}
<commit_msg>Remove unused parts of PasteInfo<commit_after>/* Copyright (c) 2014, Daniel Martí <mvdan@mvdan.cc> */
/* See LICENSE for licensing information */

package main

import (
	"crypto/rand"
	"encoding/hex"
	"errors"
	"flag"
	"fmt"
	"html/template"
	"io"
	"log"
	"net/http"
	"os"
	"path"
	"path/filepath"
	"regexp"
	"strconv"
	"strings"
	"sync"
	"time"
)

const (
	indexTmpl = "index.html"
	formTmpl  = "form.html"
	idSize    = 8
	randTries = 10

	// GET error messages
	invalidId     = "Invalid paste id."
	pasteNotFound = "Paste doesn't exist."
	unknownError  = "Something went terribly wrong."
	// POST error messages
	missingForm = "Paste could not be found inside the posted form."
)

var (
	siteUrl, listen, dataDir string
	lifeTime                 time.Duration
	maxSizeStr               string
	maxSize                  ByteSize

	regexByteSize = regexp.MustCompile(`^([\d\.]+)\s*([KM]?B|[BKM])$`)
	indexTemplate *template.Template
	formTemplate  *template.Template
	data          = struct {
		sync.RWMutex
		m map[Id]PasteInfo
	}{m: make(map[Id]PasteInfo)}
)

func init() {
	flag.StringVar(&siteUrl, "u", "http://localhost:8080", "URL of the site")
	flag.StringVar(&listen, "l", "localhost:8080", "Host and port to listen to")
	flag.StringVar(&dataDir, "d", "data", "Directory to store all the pastes in")
	flag.DurationVar(&lifeTime, "t", 12*time.Hour, "Lifetime of the pastes (units: s,m,h)")
	flag.StringVar(&maxSizeStr, "s", "1M", "Maximum size of POSTs in bytes (units: B,K,M)")
}

type PasteInfo struct {
	ModTime     time.Time
	Etag        string
	ContentType string
}

type Id [idSize / 2]byte

func IdFromString(hexId string) (Id, error) {
	var id Id
	if len(hexId) != idSize {
		return id, errors.New("Invalid id")
	}
	b, err := hex.DecodeString(hexId)
	if err != nil || len(b) != idSize/2 {
		return id, errors.New("Invalid id")
	}
	copy(id[:], b)
	return id, nil
}

func IdFromPath(idPath string) (Id, error) {
	var id Id
	parts := strings.Split(idPath, string(filepath.Separator))
	if len(parts) != 2 {
		return id, errors.New("Found invalid number of directories at " + idPath)
	}
	return IdFromString(parts[0] + parts[1])
}

func RandomId() (Id, error) {
	var id Id
	data.RLock()
	defer data.RUnlock()
	for try := 0; try < randTries; try++ {
		_, err := rand.Read(id[:])
		if err != nil {
			return id, err
		}
		if _, e := data.m[id]; !e {
			return id, nil
		}
	}
	return id, fmt.Errorf("Gave up trying to find an unused random id after %d tries", randTries)
}

func (id Id) String() string {
	return hex.EncodeToString(id[:])
}

func (id Id) Path() string {
	hexId := id.String()
	return path.Join(hexId[0:2], hexId[2:])
}

func (id Id) EndLife() {
	data.Lock()
	defer data.Unlock()
	err := os.Remove(id.Path())
	if err == nil {
		delete(data.m, id)
		log.Printf("Removed paste: %s", id)
	} else {
		log.Printf("Could not end the life of %s: %s", id, err)
		id.EndLifeAfter(2 * time.Minute)
	}
}

func (id Id) EndLifeAfter(duration time.Duration) {
	timer := time.NewTimer(duration)
	go func() {
		<-timer.C
		id.EndLife()
	}()
}

type ByteSize int64

const (
	B ByteSize = 1 << (10 * iota)
	KB
	MB
)

func parseByteSize(str string) (ByteSize, error) {
	if !regexByteSize.MatchString(str) {
		return 0, errors.New("Could not parse size in bytes")
	}
	parts := regexByteSize.FindStringSubmatch(str)
	size, _ := strconv.ParseFloat(string(parts[1]), 64)

	switch string(parts[2]) {
	case "KB", "K":
		size *= float64(KB)
	case "MB", "M":
		size *= float64(MB)
	}
	return ByteSize(size), nil
}

func (b ByteSize) String() string {
	switch {
	case b >= MB:
		return fmt.Sprintf("%.2f MB", float64(b)/float64(MB))
	case b >= KB:
		return fmt.Sprintf("%.2f KB", float64(b)/float64(KB))
	}
	return fmt.Sprintf("%d B", b)
}

func handler(w http.ResponseWriter, r *http.Request) {
	var err error
	switch r.Method {
	case "GET":
		switch r.URL.Path {
		case "/":
			indexTemplate.Execute(w, struct {
				SiteUrl, LifeTime string
			}{
				siteUrl,
				fmt.Sprintf("%g hours and %g minutes", lifeTime.Hours(),
					lifeTime.Minutes()-lifeTime.Hours()*60),
			})
			return
		case "/form":
			formTemplate.Execute(w, struct{ SiteUrl string }{siteUrl})
			return
		}
		id, err := IdFromString(r.URL.Path[1:])
		if err != nil {
			w.WriteHeader(http.StatusBadRequest)
			fmt.Fprintf(w, "%s\n", invalidId)
			return
		}
		data.RLock()
		defer data.RUnlock()
		pasteInfo, e := data.m[id]
		if !e {
			w.WriteHeader(http.StatusNotFound)
			fmt.Fprintf(w, "%s\n", pasteNotFound)
			return
		}
		if inm := r.Header.Get("If-None-Match"); inm != "" {
			if pasteInfo.Etag == inm || inm == "*" {
				w.WriteHeader(http.StatusNotModified)
				return
			}
		}
		pastePath := id.Path()
		pasteFile, err := os.Open(pastePath)
		if err != nil {
			w.WriteHeader(http.StatusInternalServerError)
			fmt.Fprintf(w, "%s\n", unknownError)
			return
		}
		defer pasteFile.Close()
		w.Header().Set("Etag", pasteInfo.Etag)
		w.Header().Set("Content-Type", pasteInfo.ContentType)
		http.ServeContent(w, r, "", pasteInfo.ModTime, pasteFile)

	case "POST":
		r.Body = http.MaxBytesReader(w, r.Body, int64(maxSize))
		var content string
		if err = r.ParseMultipartForm(int64(maxSize)); err != nil {
			log.Printf("Could not parse POST multipart form: %s", err)
			w.WriteHeader(http.StatusBadRequest)
			fmt.Fprintf(w, "%s\n", err)
			return
		}
		if vs, found := r.Form["paste"]; found && len(vs[0]) > 0 {
			content = vs[0]
		} else {
			w.WriteHeader(http.StatusBadRequest)
			fmt.Fprintf(w, "%s\n", missingForm)
			return
		}
		id, err := RandomId()
		if err == nil {
			data.Lock()
			defer data.Unlock()
		} else {
			log.Println(err)
			w.WriteHeader(http.StatusInternalServerError)
			fmt.Fprintf(w, "%s\n", unknownError)
			return
		}
		pastePath := id.Path()
		dir, _ := path.Split(pastePath)
		if err = os.Mkdir(dir, 0700); err != nil {
			log.Printf("Could not create directories leading to %s: %s", pastePath, err)
			w.WriteHeader(http.StatusInternalServerError)
			fmt.Fprintf(w, "%s\n", unknownError)
			return
		}
		id.EndLifeAfter(lifeTime)
		pasteFile, err := os.OpenFile(pastePath, os.O_WRONLY|os.O_CREATE|os.O_EXCL, 0600)
		if err != nil {
			log.Printf("Could not create new paste file %s: %s", pastePath, err)
			w.WriteHeader(http.StatusInternalServerError)
			fmt.Fprintf(w, "%s\n", unknownError)
			return
		}
		defer pasteFile.Close()
		b, err := io.WriteString(pasteFile, content)
		if err != nil {
			log.Printf("Could not write data into %s: %s", pastePath, err)
			w.WriteHeader(http.StatusInternalServerError)
			fmt.Fprintf(w, "%s\n", unknownError)
			return
		}
		pasteInfo := PasteInfo{
			ModTime:   time.Now(),
		}
		pasteInfo.Etag = fmt.Sprintf("%d-%s", pasteInfo.ModTime.Unix(), id)
		pasteInfo.ContentType = http.DetectContentType([]byte(content))
		if pasteInfo.ContentType == "application/octet-stream" {
			pasteInfo.ContentType = "text-plain; charset=utf-8"
		}
		data.m[id] = pasteInfo
		log.Printf("Created new paste %s (%s %s)", id, pasteInfo.ContentType, ByteSize(b))
		fmt.Fprintf(w, "%s/%s\n", siteUrl, id)
	}
}

func walkFunc(filePath string, fileInfo os.FileInfo, err error) error {
	if err != nil {
		return err
	}
	if fileInfo.IsDir() {
		return nil
	}
	id, err := IdFromPath(filePath)
	if err != nil {
		return errors.New("Found incompatible id at path " + filePath)
	}
	modTime := fileInfo.ModTime()
	deathTime := modTime.Add(lifeTime)
	now := time.Now()
	if deathTime.Before(now) {
		go id.EndLife()
		return nil
	}
	var lifeLeft time.Duration
	if deathTime.After(now.Add(lifeTime)) {
		lifeLeft = lifeTime
	} else {
		lifeLeft = deathTime.Sub(now)
	}
	pasteFile, err := os.Open(filePath)
	if err != nil {
		return err
	}
	defer pasteFile.Close()
	read := make([]byte, 512)
	_, err = pasteFile.Read(read)
	if err != nil {
		return err
	}
	pasteInfo := PasteInfo{
		ModTime:   modTime,
		Etag:      fmt.Sprintf("%d-%s", modTime.Unix(), id),
	}
	pasteInfo.ContentType = http.DetectContentType(read)
	if pasteInfo.ContentType == "application/octet-stream" {
		pasteInfo.ContentType = "text-plain; charset=utf-8"
	}
	data.m[id] = pasteInfo
	log.Printf("Recovered paste %s (%s %s) from %s has %s left",
		id, pasteInfo.ContentType, ByteSize(fileInfo.Size()), pasteInfo.ModTime, lifeLeft)
	id.EndLifeAfter(lifeLeft)
	return nil
}

func main() {
	var err error
	flag.Parse()
	if maxSize, err = parseByteSize(maxSizeStr); err != nil {
		log.Fatalf("Invalid max size '%s': %s", maxSizeStr, err)
	}
	if indexTemplate, err = template.ParseFiles(indexTmpl); err != nil {
		log.Fatalf("Could not load template %s: %s", indexTmpl, err)
	}
	if formTemplate, err = template.ParseFiles(formTmpl); err != nil {
		log.Fatalf("Could not load template %s: %s", formTmpl, err)
	}
	if err = os.MkdirAll(dataDir, 0700); err != nil {
		log.Fatalf("Could not create data directory %s: %s", dataDir, err)
	}
	if err = os.Chdir(dataDir); err != nil {
		log.Fatalf("Could not enter data directory %s: %s", dataDir, err)
	}
	log.Printf("maxSize  = %s", maxSize)
	log.Printf("siteUrl  = %s", siteUrl)
	log.Printf("listen   = %s", listen)
	log.Printf("dataDir  = %s", dataDir)
	log.Printf("lifeTime = %s", lifeTime)
	if err = filepath.Walk(".", walkFunc); err != nil {
		log.Fatalf("Could not recover data directory %s: %s", dataDir, err)
	}
	http.HandleFunc("/", handler)
	log.Printf("Up and running!")
	log.Fatal(http.ListenAndServe(listen, nil))
}
<|endoftext|>