<commit_before>package sdees

import (
	"hash/fnv"
	"math/rand"
	"strconv"
	"strings"
)

// ShortEncrypt runs a XOR encryption on the input string using the random bytes
// in the massive key.
// Random bytes are used starting at a position based on the hash of the input string.
// The starting position is saved as a prefix to the encrypted string
func ShortEncrypt(input string) string {
	if strings.Contains(input, ".sdz") || len(input) == 0 {
		return input
	}
	key := Cryptkey
	inputb := []byte(input)

	// Get integer hash of input, using some of the random bytes in key as salt
	h := fnv.New32a()
	h.Write(append(inputb, []byte(key)[:100]...))
	inputToNum := h.Sum32()

	// Use random integer to seed and generate random start position
	rand.Seed(int64(inputToNum))
	startPos := rand.Intn(9999000-1) + 1

	// Do XOR encryption based on that start position
	keyb := []byte(key[startPos : startPos+len(input)])
	b := make([]byte, len(inputb))
	for i := 0; i < len(inputb); i++ {
		b[i] = inputb[i] ^ keyb[i]
	}

	// Return string as [startposition]-==-[encryptedstring]
	startPosString := strconv.Itoa(startPos)
	return startPosString + "." + EncodeToString(b) + ".sdz"
}

// ShortDecrypt runs a XOR encryption on the input string using the random bytes
// in the massive key.
// Random bytes are used starting at a position based on the prefix in the input
func ShortDecrypt(input string) string {
	if !strings.Contains(input, ".sdz") {
		return input
	}
	key := Cryptkey
	parts := strings.Split(input, ".")
	inputb := DecodeString(parts[1])
	startPos, _ := strconv.Atoi(parts[0])
	keyb := []byte(key[startPos : startPos+len(inputb)])
	b := make([]byte, len(inputb))
	for i := 0; i < len(inputb); i++ {
		b[i] = inputb[i] ^ keyb[i]
	}
	return string(b)
}
<commit_msg>sdz -> otp<commit_after>package sdees

import (
	"hash/fnv"
	"math/rand"
	"strconv"
	"strings"
)

// ShortEncrypt runs a XOR encryption on the input string using the random bytes
// in the massive key.
// Random bytes are used starting at a position based on the hash of the input string.
// The starting position is saved as a prefix to the encrypted string
func ShortEncrypt(input string) string {
	if strings.Contains(input, ".otp") || len(input) == 0 {
		return input
	}
	key := Cryptkey
	inputb := []byte(input)

	// Get integer hash of input, using some of the random bytes in key as salt
	h := fnv.New32a()
	h.Write(append(inputb, []byte(key)[:100]...))
	inputToNum := h.Sum32()

	// Use random integer to seed and generate random start position
	rand.Seed(int64(inputToNum))
	startPos := rand.Intn(9999000-1) + 1

	// Do XOR encryption based on that start position
	keyb := []byte(key[startPos : startPos+len(input)])
	b := make([]byte, len(inputb))
	for i := 0; i < len(inputb); i++ {
		b[i] = inputb[i] ^ keyb[i]
	}

	// Return string as [startposition]-==-[encryptedstring]
	startPosString := strconv.Itoa(startPos)
	return startPosString + "." + EncodeToString(b) + ".otp"
}

// ShortDecrypt runs a XOR encryption on the input string using the random bytes
// in the massive key.
// Random bytes are used starting at a position based on the prefix in the input
func ShortDecrypt(input string) string {
	if !strings.Contains(input, ".otp") {
		return input
	}
	key := Cryptkey
	parts := strings.Split(input, ".")
	inputb := DecodeString(parts[1])
	startPos, _ := strconv.Atoi(parts[0])
	keyb := []byte(key[startPos : startPos+len(inputb)])
	b := make([]byte, len(inputb))
	for i := 0; i < len(inputb); i++ {
		b[i] = inputb[i] ^ keyb[i]
	}
	return string(b)
}
<|endoftext|>


<commit_before>package engine

import (
	"fmt"
	"protocol"
)

func getJoinYield(table1, table2 string, yield func(*protocol.Series) error) func(*protocol.Series) error {
	var lastPoint1 *protocol.Point
	var lastPoint2 *protocol.Point
	name := table1 + "_join_" + table2

	return mergeYield(table1, table2, func(s *protocol.Series) error {
		if *s.Name == table1 {
			lastPoint1 = s.Points[len(s.Points)-1]
			fmt.Printf("values 1: %#v\n", lastPoint1.Values)
		}

		if *s.Name == table2 {
			lastPoint2 = s.Points[len(s.Points)-1]
			fmt.Printf("values 2: %#v\n", lastPoint2.Values)
		}

		if lastPoint1 == nil || lastPoint2 == nil {
			return nil
		}

		// join the last two points and yield them
		newValues := []*protocol.FieldValue{}
		for idx, value := range lastPoint1.Values {
			if value == nil {
				value = lastPoint2.Values[idx]
			}
			newValues = append(newValues, value)
		}

		newSeries := &protocol.Series{
			Name:   &name,
			Fields: s.Fields,
			Points: []*protocol.Point{
				&protocol.Point{
					Values:         newValues,
					Timestamp:      lastPoint2.Timestamp,
					SequenceNumber: lastPoint2.SequenceNumber,
				},
			},
		}

		lastPoint1 = nil
		lastPoint2 = nil

		return yield(newSeries)
	})
}

func getMergeYield(table1, table2 string, yield func(*protocol.Series) error) func(*protocol.Series) error {
	name := table1 + "_merge_" + table2

	return mergeYield(table1, table2, func(s *protocol.Series) error {
		s.Name = &name
		return yield(s)
	})
}

type MergeState struct {
	name       string
	series     []*protocol.Series
	fields     []string
	nullValues []*protocol.FieldValue
	isLeft     bool
	done       bool
}

func (self *MergeState) hasPoints() bool {
	return len(self.series) > 0 && len(self.series[0].Points) > 0
}

func (self *MergeState) isEarlier(other *MergeState) bool {
	return *self.series[0].Points[0].Timestamp < *other.series[0].Points[0].Timestamp
}

func (self *MergeState) mergeValues(other *MergeState, p *protocol.Point) {
	if self.isLeft {
		p.Values = append(p.Values, other.nullValues...)
		return
	}
	p.Values = append(other.nullValues, p.Values...)
}

func (self *MergeState) flushIfOtherStateIsEmpty(other *MergeState, fields []string, yield func(*protocol.Series) error) error {
	if other.done && len(other.series) == 0 {
		for _, s := range self.series {
			for _, p := range s.Points {
				self.mergeValues(other, p)
			}

			err := yield(&protocol.Series{
				Name:   s.Name,
				Fields: fields,
				Points: s.Points,
			})
			if err != nil {
				return err
			}
		}
		self.series = nil
	}
	return nil
}

// update the state, the points belong to this MergeState (i.e. the name of the timeseries matches)
func (self *MergeState) updateState(p *protocol.Series) {
	if *p.Name != self.name {
		return
	}

	// setup the fields
	if self.fields == nil {
		for _, f := range p.Fields {
			self.fields = append(self.fields, self.name+"."+f)
		}
		for i := 0; i < len(p.Fields); i++ {
			self.nullValues = append(self.nullValues, nil)
		}
	}

	// data for current table is exhausted
	if len(p.Points) == 0 {
		self.done = true
	} else {
		self.series = append(self.series, p)
	}
}

// returns a yield function that will sort points from table1 and
// table2 no matter what the order in which they are received.
//
// FIXME: This won't work for queries that are executed in descending order.
func mergeYield(table1, table2 string, yield func(*protocol.Series) error) func(*protocol.Series) error {
	state1 := &MergeState{
		name:   table1,
		isLeft: true,
	}
	state2 := &MergeState{
		name:   table2,
		isLeft: false,
	}

	return func(p *protocol.Series) error {
		state1.updateState(p)
		state2.updateState(p)

		mergedFields := append(state1.fields, state2.fields...)

		// if one of the states (let's call it s1) is done (i.e. we'll
		// receive no more points for that series) then we know that we
		// won't get any points for s1 that are older than what's in s2 so
		// we can safely flush all s2's points.
		if err := state1.flushIfOtherStateIsEmpty(state2, mergedFields, yield); err != nil {
			return err
		}
		if err := state2.flushIfOtherStateIsEmpty(state1, mergedFields, yield); err != nil {
			return err
		}

		for state1.hasPoints() && state2.hasPoints() {
			var points []*protocol.Point

			state := state1
			otherState := state2
			if state2.isEarlier(state1) {
				state = state2
				otherState = state1
			}
			points = []*protocol.Point{state.series[0].Points[0]}
			for _, p := range points {
				state.mergeValues(otherState, p)
			}
			// get rid of that point, or get rid of the entire series
			// if this is the last point
			if len(state.series[0].Points) == 1 {
				state.series = state.series[1:]
			} else {
				state.series[0].Points = state.series[0].Points[1:]
			}
			err := yield(&protocol.Series{
				Name:   &state.name,
				Fields: mergedFields,
				Points: points,
			})
			if err != nil {
				return err
			}
		}
		return nil
	}
}
<commit_msg>some documentation and small refactoring.<commit_after>package engine

import (
	"fmt"
	"protocol"
)

func getJoinYield(table1, table2 string, yield func(*protocol.Series) error) func(*protocol.Series) error {
	var lastPoint1 *protocol.Point
	var lastPoint2 *protocol.Point
	name := table1 + "_join_" + table2

	return mergeYield(table1, table2, func(s *protocol.Series) error {
		if *s.Name == table1 {
			lastPoint1 = s.Points[len(s.Points)-1]
			fmt.Printf("values 1: %#v\n", lastPoint1.Values)
		}

		if *s.Name == table2 {
			lastPoint2 = s.Points[len(s.Points)-1]
			fmt.Printf("values 2: %#v\n", lastPoint2.Values)
		}

		if lastPoint1 == nil || lastPoint2 == nil {
			return nil
		}

		// join the last two points and yield them
		newValues := []*protocol.FieldValue{}
		for idx, value := range lastPoint1.Values {
			if value == nil {
				value = lastPoint2.Values[idx]
			}
			newValues = append(newValues, value)
		}

		newSeries := &protocol.Series{
			Name:   &name,
			Fields: s.Fields,
			Points: []*protocol.Point{
				&protocol.Point{
					Values:         newValues,
					Timestamp:      lastPoint2.Timestamp,
					SequenceNumber: lastPoint2.SequenceNumber,
				},
			},
		}

		lastPoint1 = nil
		lastPoint2 = nil

		return yield(newSeries)
	})
}

func getMergeYield(table1, table2 string, yield func(*protocol.Series) error) func(*protocol.Series) error {
	name := table1 + "_merge_" + table2

	return mergeYield(table1, table2, func(s *protocol.Series) error {
		s.Name = &name
		return yield(s)
	})
}

type mergeState struct {
	name       string
	series     []*protocol.Series
	fields     []string
	nullValues []*protocol.FieldValue
	isLeft     bool
	done       bool
}

func (self *mergeState) hasPoints() bool {
	return len(self.series) > 0 && len(self.series[0].Points) > 0
}

func (self *mergeState) isEarlier(other *mergeState) bool {
	return *self.series[0].Points[0].Timestamp < *other.series[0].Points[0].Timestamp
}

// set the fields of the other time series to null making sure that
// the order of the null values match the order of the field
// definitions, i.e. left timeseries first followed by values from the
// right timeseries
func (self *mergeState) mergeValues(other *mergeState, p *protocol.Point) {
	if self.isLeft {
		p.Values = append(p.Values, other.nullValues...)
		return
	}
	p.Values = append(other.nullValues, p.Values...)
}

// if `other` state is done (i.e. we'll receive no more points for its
// timeseries) then we know that we won't get any points that are
// older than what's in `self` so we can safely flush all `self`
// points.
func (self *mergeState) flushIfOtherStateIsEmpty(other *mergeState, fields []string, yield func(*protocol.Series) error) error {
	if other.done && len(other.series) == 0 {
		for _, s := range self.series {
			for _, p := range s.Points {
				self.mergeValues(other, p)
			}

			err := yield(&protocol.Series{
				Name:   s.Name,
				Fields: fields,
				Points: s.Points,
			})
			if err != nil {
				return err
			}
		}
		self.series = nil
	}
	return nil
}

// update the state, the points belong to this mergeState (i.e. the name of the timeseries matches)
func (self *mergeState) updateState(p *protocol.Series) {
	if *p.Name != self.name {
		return
	}

	// setup the fields
	if self.fields == nil {
		for _, f := range p.Fields {
			self.fields = append(self.fields, self.name+"."+f)
		}
		for i := 0; i < len(p.Fields); i++ {
			self.nullValues = append(self.nullValues, nil)
		}
	}

	// data for current table is exhausted
	if len(p.Points) == 0 {
		self.done = true
	} else {
		self.series = append(self.series, p)
	}
}

func (self *mergeState) removeAndGetFirstPoint() *protocol.Point {
	point := self.series[0].Points[0]
	// get rid of that point, or get rid of the entire series
	// if this is the last point
	if len(self.series[0].Points) == 1 {
		self.series = self.series[1:]
	} else {
		self.series[0].Points = self.series[0].Points[1:]
	}
	return point
}

// returns a yield function that will sort points from table1 and
// table2 no matter what the order in which they are received.
//
// FIXME: This won't work for queries that are executed in descending order.
func mergeYield(table1, table2 string, yield func(*protocol.Series) error) func(*protocol.Series) error {
	state1 := &mergeState{
		name:   table1,
		isLeft: true,
	}
	state2 := &mergeState{
		name:   table2,
		isLeft: false,
	}

	return func(p *protocol.Series) error {
		state1.updateState(p)
		state2.updateState(p)

		mergedFields := append(state1.fields, state2.fields...)

		if err := state1.flushIfOtherStateIsEmpty(state2, mergedFields, yield); err != nil {
			return err
		}
		if err := state2.flushIfOtherStateIsEmpty(state1, mergedFields, yield); err != nil {
			return err
		}

		// see which point should be returned next and remove it from the
		// series
		for state1.hasPoints() && state2.hasPoints() {
			// state is the state of the series from which the next point
			// will be fetched
			state := state1
			otherState := state2
			if state2.isEarlier(state1) {
				state = state2
				otherState = state1
			}

			p := state.removeAndGetFirstPoint()
			state.mergeValues(otherState, p)

			err := yield(&protocol.Series{
				Name:   &state.name,
				Fields: mergedFields,
				Points: []*protocol.Point{p},
			})
			if err != nil {
				return err
			}
		}
		return nil
	}
}
<|endoftext|>


<commit_before>package engine

import (
    "cgl.tideland.biz/applog"
    "time"
    "vector"
)

var (
    PLAYER_UPDATE_TIMER, _ = time.ParseDuration("50ms")
)

// TODO: Make this a non-global variable
var world = MakeWorld(10, 10)

type WorldUpdateFunc func (*World)

type World struct {
    players map[string] *Player
    worldMap *WorldMap
    currentId uint

    register chan *connection
    unregister chan *Client
    update chan WorldUpdateFunc
}

func MakeWorld(width, height uint) *World {
    world := &World {
        make(map [string] *Player),
        MakeWorldMap(width, height),
        0,
        make(chan *connection),
        make(chan *Client),
        make(chan WorldUpdateFunc),
    }
    gm := MakeGameChunk(width, height, vector.Vector2{0, 0})
    world.worldMap.SetGameChunk(vector.Vector2{0, 0}, gm)
    return world
}

func (w *World) StartPlayerLogin(name, token string) bool {
    p, found := w.GetPlayer(name)
    if found {
        p.Login(token)
        return true
    } else {
        return false
    }
}

func (w *World) GetPlayer(name string) (*Player, bool) {
    p := w.players[name]
    if p != nil {
        return p, true
    } else {
        return nil, false
    }
}

func (w *World) NewPlayer(name string, token string) *Player {
    path := "players/" + name
    if res, _ := fileSystem.FileExists(path); res {
        p, err := fileSystem.LoadPlayer(path)
        if err != nil {
            applog.Errorf("Could not load player at path: %s. Error: %s",
                          path, err)
            return nil
        }
        w.players[p.Name] = p
        return p
    } else {
        applog.Debugf("Creating a new player")
        p := MakePlayer(w.getNextId(), name, vector.Vector2{0, 0})
        w.players[p.Name] = p
        return p
    }
    return nil
}

func (w *World) Start() {
    go w.HandleConnections()
    go w.UpdateLoop()
    go w.HandleShutdown()
}

func (w *World) SendPlayerInitialInfo(p *Player) {
    p.write(MakeMapPacket(w.worldMap.grid))
    p.write(MakePlayerPacket(p))
    RepeatingTimer(PLAYER_UPDATE_TIMER, func () bool {
        p.write(MakeNearbyPlayerUpdatePacket(w))
        return p.IsOnline()
    })
}

func (w *World) HandleConnections() {
    for conn := range w.register {
        applog.Debugf("Received request to register new player")
        client := MakeClient()
        go conn.writePump(client)
        go conn.readPump(client)
        go client.read()
    }
}

func (w *World) HandleShutdown() {
    for client := range w.unregister {
        applog.Debugf("Client %s logging out", client)
        path := "players/" + client.player.Name
        err := fileSystem.Save(path, client.player)
        if err != nil {
            applog.Errorf("Failed to save player at path: %s. Error: %s",
                          path, err)
        }
    }
}

func (w *World) Broadcast(payload GameWriter) {
    for _, player := range w.players {
        player.write(payload)
    }
}

func (w *World) UpdateLoop() {
    for updateFunction := range w.update {
        updateFunction(w)
    }
}

func (w *World) getNextId() uint {
    id := w.currentId
    w.currentId++
    return id
}
<commit_msg>Fix server crash<commit_after>package engine

import (
    "cgl.tideland.biz/applog"
    "time"
    "vector"
)

var (
    PLAYER_UPDATE_TIMER, _ = time.ParseDuration("50ms")
)

// TODO: Make this a non-global variable
var world = MakeWorld(10, 10)

type WorldUpdateFunc func (*World)

type World struct {
    players map[string] *Player
    worldMap *WorldMap
    currentId uint

    register chan *connection
    unregister chan *Client
    update chan WorldUpdateFunc
}

func MakeWorld(width, height uint) *World {
    world := &World {
        make(map [string] *Player),
        MakeWorldMap(width, height),
        0,
        make(chan *connection),
        make(chan *Client),
        make(chan WorldUpdateFunc),
    }
    gm := MakeGameChunk(width, height, vector.Vector2{0, 0})
    world.worldMap.SetGameChunk(vector.Vector2{0, 0}, gm)
    return world
}

func (w *World) StartPlayerLogin(name, token string) bool {
    p, found := w.GetPlayer(name)
    if found {
        p.Login(token)
        return true
    } else {
        return false
    }
}

func (w *World) GetPlayer(name string) (*Player, bool) {
    p := w.players[name]
    if p != nil {
        return p, true
    } else {
        return nil, false
    }
}

func (w *World) NewPlayer(name string, token string) *Player {
    path := "players/" + name
    if res, _ := fileSystem.FileExists(path); res {
        p, err := fileSystem.LoadPlayer(path)
        if err != nil {
            applog.Errorf("Could not load player at path: %s. Error: %s",
                          path, err)
            return nil
        }
        w.players[p.Name] = p
        return p
    } else {
        applog.Debugf("Creating a new player")
        p := MakePlayer(w.getNextId(), name, vector.Vector2{0, 0})
        w.players[p.Name] = p
        return p
    }
    return nil
}

func (w *World) Start() {
    go w.HandleConnections()
    go w.UpdateLoop()
    go w.HandleShutdown()
}

func (w *World) SendPlayerInitialInfo(p *Player) {
    p.write(MakeMapPacket(w.worldMap.grid))
    p.write(MakePlayerPacket(p))
    RepeatingTimer(PLAYER_UPDATE_TIMER, func () bool {
        p.write(MakeNearbyPlayerUpdatePacket(w))
        return p.IsOnline()
    })
}

func (w *World) HandleConnections() {
    for conn := range w.register {
        applog.Debugf("Received request to register new player")
        client := MakeClient()
        go conn.writePump(client)
        go conn.readPump(client)
        go client.read()
    }
}

func (w *World) HandleShutdown() {
    for client := range w.unregister {
        applog.Debugf("Client %s logging out", client)
        if client.player == nil {
            continue
        }
        path := "players/" + client.player.Name
        err := fileSystem.Save(path, client.player)
        if err != nil {
            applog.Errorf("Failed to save player at path: %s. Error: %s",
                          path, err)
        }
    }
}

func (w *World) Broadcast(payload GameWriter) {
    for _, player := range w.players {
        player.write(payload)
    }
}

func (w *World) UpdateLoop() {
    for updateFunction := range w.update {
        updateFunction(w)
    }
}

func (w *World) getNextId() uint {
    id := w.currentId
    w.currentId++
    return id
}
<|endoftext|>


<commit_before>package main

import (
	"../util"
	"fmt"
	"io"
	"log"
	"net/http"
	"sort"
	"time"
)

type ClientStats struct {
	version       string
	name          string
	state         int
	inFlightCount int64
	readyCount    int64
	messageCount  uint64
	finishCount   uint64
	requeueCount  uint64
	connectTime   time.Time
}

type Topics []*Topic
type TopicsByName struct {
	Topics
}
type Channels []*Channel
type ChannelsByName struct {
	Channels
}

func (t Topics) Len() int                   { return len(t) }
func (t Topics) Swap(i, j int)              { t[i], t[j] = t[j], t[i] }
func (c Channels) Len() int                 { return len(c) }
func (c Channels) Swap(i, j int)            { c[i], c[j] = c[j], c[i] }
func (t TopicsByName) Less(i, j int) bool   { return t.Topics[i].name < t.Topics[j].name }
func (c ChannelsByName) Less(i, j int) bool { return c.Channels[i].name < c.Channels[j].name }

// print out stats for each topic/channel
func statsHandler(w http.ResponseWriter, req *http.Request) {
	reqParams, err := util.NewReqParams(req)
	if err != nil {
		log.Printf("ERROR: failed to parse request params - %s", err.Error())
		w.Write(util.ApiResponse(500, "INVALID_REQUEST", nil))
		return
	}

	formatString, _ := reqParams.Query("format")
	jsonFormat := formatString == "json"
	now := time.Now()

	nsqd.RLock()
	defer nsqd.RUnlock()

	if len(nsqd.topicMap) == 0 {
		if jsonFormat {
			w.Write(util.ApiResponse(500, "NO_TOPICS", nil))
		} else {
			io.WriteString(w, "NO_TOPICS\n")
		}
		return
	}

	realTopics := make([]*Topic, len(nsqd.topicMap))
	topics := make([]interface{}, len(nsqd.topicMap))
	topic_index := 0
	for _, t := range nsqd.topicMap {
		realTopics[topic_index] = t
		topic_index++
	}

	sort.Sort(TopicsByName{realTopics})
	for topic_index, t := range realTopics {
		t.RLock()

		if !jsonFormat {
			io.WriteString(w, fmt.Sprintf("\n[%s] depth: %-5d be-depth: %-5d msgs %-8d\n",
				t.name,
				int64(len(t.memoryMsgChan))+t.backend.Depth(),
				t.backend.Depth(),
				t.messageCount))
		}

		realChannels := make([]*Channel, len(t.channelMap))
		channel_index := 0
		for _, c := range t.channelMap {
			realChannels[channel_index] = c
			channel_index++
		}
		sort.Sort(ChannelsByName{realChannels})

		channels := make([]interface{}, len(t.channelMap))
		for channel_index, c := range realChannels {
			c.RLock()
			if jsonFormat {
				clients := make([]interface{}, len(c.clients))
				for client_index, client := range c.clients {
					clientStats := client.Stats()
					clients[client_index] = struct {
						Version       string `json:"version"`
						Name          string `json:"name"`
						State         int    `json:"state"`
						ReadyCount    int64  `json:"ready_count"`
						InFlightCount int64  `json:"in_flight_count"`
						MessageCount  uint64 `json:"message_count"`
						FinishCount   uint64 `json:"finish_count"`
						RequeueCount  uint64 `json:"requeue_count"`
						ConnectTime   int64  `json:"connect_ts"`
					}{
						clientStats.version,
						clientStats.name,
						clientStats.state,
						clientStats.readyCount,
						clientStats.inFlightCount,
						clientStats.messageCount,
						clientStats.finishCount,
						clientStats.requeueCount,
						clientStats.connectTime.Unix(),
					}
				}
				channels[channel_index] = struct {
					ChannelName   string        `json:"channel_name"`
					Depth         int64         `json:"depth"`
					BackendDepth  int64         `json:"backend_depth"`
					InFlightCount int           `json:"in_flight_count"`
					DeferredCount int           `json:"deferred_count"`
					MessageCount  uint64        `json:"message_count"`
					RequeueCount  uint64        `json:"requeue_count"`
					TimeoutCount  uint64        `json:"timeout_count"`
					Clients       []interface{} `json:"clients"`
				}{
					c.name,
					int64(len(c.memoryMsgChan)) + c.backend.Depth(),
					c.backend.Depth(),
					len(c.inFlightMessages),
					len(c.deferredMessages),
					c.messageCount,
					c.requeueCount,
					c.timeoutCount,
					clients,
				}
				channel_index++
			} else {
				io.WriteString(w,
					fmt.Sprintf("    [%-10s] depth: %-5d be-depth: %-5d inflt: %-4d def: %-4d msgs: %-8d re-q: %-5d timeout: %-5d\n",
						c.name,
						int64(len(c.memoryMsgChan))+c.backend.Depth(),
						c.backend.Depth(),
						len(c.inFlightMessages),
						len(c.deferredMessages),
						c.messageCount,
						c.requeueCount,
						c.timeoutCount))
				for _, client := range c.clients {
					clientStats := client.Stats()
					duration := now.Sub(clientStats.connectTime).Seconds()
					io.WriteString(w, fmt.Sprintf("        [%s %-12s] state: %d inflt: %-4d rdy: %-4d fin: %-8d re-q: %-8d msgs: %-8d connected: %s\n",
						clientStats.version,
						clientStats.name,
						clientStats.state,
						clientStats.inFlightCount,
						clientStats.readyCount,
						clientStats.finishCount,
						clientStats.requeueCount,
						clientStats.messageCount,
						time.Duration(int64(duration))*time.Second, // truncate to the second
					))
				}
			}
			c.RUnlock()
		}

		topics[topic_index] = struct {
			TopicName    string        `json:"topic_name"`
			Channels     []interface{} `json:"channels"`
			Depth        int64         `json:"depth"`
			BackendDepth int64         `json:"backend_depth"`
			MessageCount uint64        `json:"message_count"`
		}{
			TopicName:    t.name,
			Channels:     channels,
			Depth:        int64(len(t.memoryMsgChan)) + t.backend.Depth(),
			BackendDepth: t.backend.Depth(),
			MessageCount: t.messageCount,
		}
		topic_index++

		t.RUnlock()
	}

	if jsonFormat {
		w.Write(util.ApiResponse(200, "OK", struct {
			Topics []interface{} `json:"topics"`
		}{topics}))
	}

}
<commit_msg>typo in stats output<commit_after>package main

import (
	"../util"
	"fmt"
	"io"
	"log"
	"net/http"
	"sort"
	"time"
)

type ClientStats struct {
	version       string
	name          string
	state         int
	inFlightCount int64
	readyCount    int64
	messageCount  uint64
	finishCount   uint64
	requeueCount  uint64
	connectTime   time.Time
}

type Topics []*Topic
type TopicsByName struct {
	Topics
}
type Channels []*Channel
type ChannelsByName struct {
	Channels
}

func (t Topics) Len() int                   { return len(t) }
func (t Topics) Swap(i, j int)              { t[i], t[j] = t[j], t[i] }
func (c Channels) Len() int                 { return len(c) }
func (c Channels) Swap(i, j int)            { c[i], c[j] = c[j], c[i] }
func (t TopicsByName) Less(i, j int) bool   { return t.Topics[i].name < t.Topics[j].name }
func (c ChannelsByName) Less(i, j int) bool { return c.Channels[i].name < c.Channels[j].name }

// print out stats for each topic/channel
func statsHandler(w http.ResponseWriter, req *http.Request) {
	reqParams, err := util.NewReqParams(req)
	if err != nil {
		log.Printf("ERROR: failed to parse request params - %s", err.Error())
		w.Write(util.ApiResponse(500, "INVALID_REQUEST", nil))
		return
	}

	formatString, _ := reqParams.Query("format")
	jsonFormat := formatString == "json"
	now := time.Now()

	nsqd.RLock()
	defer nsqd.RUnlock()

	if len(nsqd.topicMap) == 0 {
		if jsonFormat {
			w.Write(util.ApiResponse(500, "NO_TOPICS", nil))
		} else {
			io.WriteString(w, "NO_TOPICS\n")
		}
		return
	}

	realTopics := make([]*Topic, len(nsqd.topicMap))
	topics := make([]interface{}, len(nsqd.topicMap))
	topic_index := 0
	for _, t := range nsqd.topicMap {
		realTopics[topic_index] = t
		topic_index++
	}

	sort.Sort(TopicsByName{realTopics})
	for topic_index, t := range realTopics {
		t.RLock()

		if !jsonFormat {
			io.WriteString(w, fmt.Sprintf("\n[%s] depth: %-5d be-depth: %-5d msgs: %-8d\n",
				t.name,
				int64(len(t.memoryMsgChan))+t.backend.Depth(),
				t.backend.Depth(),
				t.messageCount))
		}

		realChannels := make([]*Channel, len(t.channelMap))
		channel_index := 0
		for _, c := range t.channelMap {
			realChannels[channel_index] = c
			channel_index++
		}
		sort.Sort(ChannelsByName{realChannels})

		channels := make([]interface{}, len(t.channelMap))
		for channel_index, c := range realChannels {
			c.RLock()
			if jsonFormat {
				clients := make([]interface{}, len(c.clients))
				for client_index, client := range c.clients {
					clientStats := client.Stats()
					clients[client_index] = struct {
						Version       string `json:"version"`
						Name          string `json:"name"`
						State         int    `json:"state"`
						ReadyCount    int64  `json:"ready_count"`
						InFlightCount int64  `json:"in_flight_count"`
						MessageCount  uint64 `json:"message_count"`
						FinishCount   uint64 `json:"finish_count"`
						RequeueCount  uint64 `json:"requeue_count"`
						ConnectTime   int64  `json:"connect_ts"`
					}{
						clientStats.version,
						clientStats.name,
						clientStats.state,
						clientStats.readyCount,
						clientStats.inFlightCount,
						clientStats.messageCount,
						clientStats.finishCount,
						clientStats.requeueCount,
						clientStats.connectTime.Unix(),
					}
				}
				channels[channel_index] = struct {
					ChannelName   string        `json:"channel_name"`
					Depth         int64         `json:"depth"`
					BackendDepth  int64         `json:"backend_depth"`
					InFlightCount int           `json:"in_flight_count"`
					DeferredCount int           `json:"deferred_count"`
					MessageCount  uint64        `json:"message_count"`
					RequeueCount  uint64        `json:"requeue_count"`
					TimeoutCount  uint64        `json:"timeout_count"`
					Clients       []interface{} `json:"clients"`
				}{
					c.name,
					int64(len(c.memoryMsgChan)) + c.backend.Depth(),
					c.backend.Depth(),
					len(c.inFlightMessages),
					len(c.deferredMessages),
					c.messageCount,
					c.requeueCount,
					c.timeoutCount,
					clients,
				}
				channel_index++
			} else {
				io.WriteString(w,
					fmt.Sprintf("    [%-10s] depth: %-5d be-depth: %-5d inflt: %-4d def: %-4d msgs: %-8d re-q: %-5d timeout: %-5d\n",
						c.name,
						int64(len(c.memoryMsgChan))+c.backend.Depth(),
						c.backend.Depth(),
						len(c.inFlightMessages),
						len(c.deferredMessages),
						c.messageCount,
						c.requeueCount,
						c.timeoutCount))
				for _, client := range c.clients {
					clientStats := client.Stats()
					duration := now.Sub(clientStats.connectTime).Seconds()
					io.WriteString(w, fmt.Sprintf("        [%s %-12s] state: %d inflt: %-4d rdy: %-4d fin: %-8d re-q: %-8d msgs: %-8d connected: %s\n",
						clientStats.version,
						clientStats.name,
						clientStats.state,
						clientStats.inFlightCount,
						clientStats.readyCount,
						clientStats.finishCount,
						clientStats.requeueCount,
						clientStats.messageCount,
						time.Duration(int64(duration))*time.Second, // truncate to the second
					))
				}
			}
			c.RUnlock()
		}

		topics[topic_index] = struct {
			TopicName    string        `json:"topic_name"`
			Channels     []interface{} `json:"channels"`
			Depth        int64         `json:"depth"`
			BackendDepth int64         `json:"backend_depth"`
			MessageCount uint64        `json:"message_count"`
		}{
			TopicName:    t.name,
			Channels:     channels,
			Depth:        int64(len(t.memoryMsgChan)) + t.backend.Depth(),
			BackendDepth: t.backend.Depth(),
			MessageCount: t.messageCount,
		}
		topic_index++

		t.RUnlock()
	}

	if jsonFormat {
		w.Write(util.ApiResponse(200, "OK", struct {
			Topics []interface{} `json:"topics"`
		}{topics}))
	}

}
<|endoftext|>


<commit_before>package v7

import (
	"fmt"
	"strings"

	"code.cloudfoundry.org/cli/actor/sharedaction"
	"code.cloudfoundry.org/cli/actor/v7action"
	"code.cloudfoundry.org/cli/command"
	"code.cloudfoundry.org/cli/command/flag"
	"code.cloudfoundry.org/cli/command/translatableerror"
	"code.cloudfoundry.org/cli/command/v7/shared"
	"code.cloudfoundry.org/cli/types"
	"code.cloudfoundry.org/clock"
)

//go:generate counterfeiter . SetLabelActor

type SetLabelActor interface {
	UpdateApplicationLabelsByApplicationName(string, string, map[string]types.NullString) (v7action.Warnings, error)
	UpdateBuildpackLabelsByBuildpackNameAndStack(string, string, map[string]types.NullString) (v7action.Warnings, error)
	UpdateDomainLabelsByDomainName(string, map[string]types.NullString) (v7action.Warnings, error)
	UpdateOrganizationLabelsByOrganizationName(string, map[string]types.NullString) (v7action.Warnings, error)
	UpdateRouteLabels(string, string, map[string]types.NullString) (v7action.Warnings, error)
	UpdateSpaceLabelsBySpaceName(string, string, map[string]types.NullString) (v7action.Warnings, error)
	UpdateStackLabelsByStackName(string, map[string]types.NullString) (v7action.Warnings, error)
	UpdateServiceBrokerLabelsByServiceBrokerName(string, map[string]types.NullString) (v7action.Warnings, error)
}

type SetLabelCommand struct {
	RequiredArgs    flag.SetLabelArgs `positional-args:"yes"`
	usage           interface{}       `usage:"CF_NAME set-label RESOURCE RESOURCE_NAME KEY=VALUE...\n\nEXAMPLES:\n   cf set-label app dora env=production\n   cf set-label org business pci=true public-facing=false\n   cf set-label buildpack go_buildpack go=1.12 -s cflinuxfs3\n\nRESOURCES:\n   app\n   buildpack\n   domain\n   org\n   route\n   service-broker\n   space\n   stack"`
	relatedCommands interface{}       `related_commands:"labels, unset-label"`
	BuildpackStack  string            `long:"stack" short:"s" description:"Specify stack to disambiguate buildpacks with the same name"`

	UI          command.UI
	Config      command.Config
	SharedActor command.SharedActor
	Actor       SetLabelActor
}

func (cmd *SetLabelCommand) Setup(config command.Config, ui command.UI) error {
	cmd.UI = ui
	cmd.Config = config
	cmd.SharedActor = sharedaction.NewActor(config)
	ccClient, _, err := shared.GetNewClientsAndConnectToCF(config, ui, "")
	if err != nil {
		return err
	}
	cmd.Actor = v7action.NewActor(ccClient, config, nil, nil, clock.NewClock())
	return nil
}

func (cmd SetLabelCommand) Execute(args []string) error {
	labels := make(map[string]types.NullString)
	for _, label := range cmd.RequiredArgs.Labels {
		parts := strings.SplitN(label, "=", 2)
		if len(parts) < 2 {
			return fmt.Errorf("Metadata error: no value provided for label '%s'", label)
		}
		labels[parts[0]] = types.NewNullString(parts[1])
	}

	username, err := cmd.Config.CurrentUserName()
	if err != nil {
		return err
	}

	if err = cmd.validateFlags(); err != nil {
		return err
	}

	switch cmd.canonicalResourceTypeForName() {
	case App:
		err = cmd.executeApp(username, labels)
	case Buildpack:
		err = cmd.executeBuildpack(username, labels)
	case Domain:
		err = cmd.executeDomain(username, labels)
	case Org:
		err = cmd.executeOrg(username, labels)
	case Route:
		err = cmd.executeRoute(username, labels)
	case Space:
		err = cmd.executeSpace(username, labels)
	case Stack:
		err = cmd.executeStack(username, labels)
	case ServiceBroker:
		err = cmd.executeServiceBroker(username, labels)
	default:
		err = fmt.Errorf("Unsupported resource type of '%s'", cmd.RequiredArgs.ResourceType)
	}

	if err != nil {
		return err
	}

	cmd.UI.DisplayOK()
	return nil
}

func (cmd SetLabelCommand) executeApp(username string, labels map[string]types.NullString) error {
	err := cmd.SharedActor.CheckTarget(true, true)
	if err != nil {
		return err
	}

	appName := cmd.RequiredArgs.ResourceName

	preFlavoringText := fmt.Sprintf("Setting label(s) for %s {{.ResourceName}} in org {{.OrgName}} / space {{.SpaceName}} as {{.User}}...", strings.ToLower(cmd.RequiredArgs.ResourceType))
	cmd.UI.DisplayTextWithFlavor(
		preFlavoringText,
		map[string]interface{}{
			"ResourceName": appName,
			"OrgName":      cmd.Config.TargetedOrganization().Name,
			"SpaceName":    cmd.Config.TargetedSpace().Name,
			"User":         username,
		},
	)

	warnings, err := cmd.Actor.UpdateApplicationLabelsByApplicationName(appName,
		cmd.Config.TargetedSpace().GUID,
		labels)

	cmd.UI.DisplayWarnings(warnings)

	return err
}

func (cmd SetLabelCommand) executeBuildpack(username string, labels map[string]types.NullString) error {
	err := cmd.SharedActor.CheckTarget(false, false)
	if err != nil {
		return err
	}

	preFlavoringText := fmt.Sprintf("Setting label(s) for %s {{.ResourceName}} as {{.User}}...", strings.ToLower(cmd.RequiredArgs.ResourceType))
	if cmd.BuildpackStack != "" {
		preFlavoringText = fmt.Sprintf("Setting label(s) for %s {{.ResourceName}} with stack {{.StackName}} as {{.User}}...", strings.ToLower(cmd.RequiredArgs.ResourceType))
	}

	cmd.UI.DisplayTextWithFlavor(
		preFlavoringText,
		map[string]interface{}{
			"ResourceName": cmd.RequiredArgs.ResourceName,
			"StackName":    cmd.BuildpackStack,
			"User":         username,
		},
	)

	warnings, err := cmd.Actor.UpdateBuildpackLabelsByBuildpackNameAndStack(cmd.RequiredArgs.ResourceName, cmd.BuildpackStack, labels)
	cmd.UI.DisplayWarnings(warnings)

	return err
}

func (cmd SetLabelCommand) executeOrg(username string, labels map[string]types.NullString) error {
	err := cmd.SharedActor.CheckTarget(false, false)
	if err != nil {
		return err
	}

	preFlavoringText := fmt.Sprintf("Setting label(s) for %s {{.ResourceName}} as {{.User}}...", strings.ToLower(cmd.RequiredArgs.ResourceType))
	cmd.UI.DisplayTextWithFlavor(
		preFlavoringText,
		map[string]interface{}{
			"ResourceName": cmd.RequiredArgs.ResourceName,
			"OrgName":      cmd.Config.TargetedOrganization().Name,
			"User":         username,
		},
	)

	warnings, err := cmd.Actor.UpdateOrganizationLabelsByOrganizationName(cmd.RequiredArgs.ResourceName,
		labels)
	cmd.UI.DisplayWarnings(warnings)

	return err
}

func (cmd SetLabelCommand) executeRoute(username string, labels map[string]types.NullString) error {
	err := cmd.SharedActor.CheckTarget(true, true)
	if err != nil {
		return err
	}

	preFlavoringText := fmt.Sprintf("Setting label(s) for %s {{.ResourceName}} in org {{.OrgName}} / space {{.SpaceName}} as {{.User}}...", strings.ToLower(cmd.RequiredArgs.ResourceType))
	cmd.UI.DisplayTextWithFlavor(
		preFlavoringText,
		map[string]interface{}{
			"ResourceName": cmd.RequiredArgs.ResourceName,
			"OrgName":      cmd.Config.TargetedOrganization().Name,
			"SpaceName":    cmd.Config.TargetedSpace().Name,
			"User":         username,
		},
	)

	warnings, err := cmd.Actor.UpdateRouteLabels(
		cmd.RequiredArgs.ResourceName,
		cmd.Config.TargetedSpace().GUID,
		labels,
	)
	cmd.UI.DisplayWarnings(warnings)

	return err
}

func (cmd SetLabelCommand) executeSpace(username string, labels map[string]types.NullString) error {
	err := cmd.SharedActor.CheckTarget(true, false)
	if err != nil {
		return err
	}

	spaceName := cmd.RequiredArgs.ResourceName

	preFlavoringText := fmt.Sprintf("Setting label(s) for %s {{.ResourceName}} in org {{.OrgName}} as {{.User}}...", strings.ToLower(cmd.RequiredArgs.ResourceType))
	cmd.UI.DisplayTextWithFlavor(
		preFlavoringText,
		map[string]interface{}{
			"ResourceName": spaceName,
			"OrgName":      cmd.Config.TargetedOrganization().Name,
			"User":         username,
		},
	)

	warnings, err := cmd.Actor.UpdateSpaceLabelsBySpaceName(spaceName,
		cmd.Config.TargetedOrganization().GUID,
		labels)

	cmd.UI.DisplayWarnings(warnings)

	return err
}

func (cmd SetLabelCommand) executeStack(username string, labels map[string]types.NullString) error {
	err := cmd.SharedActor.CheckTarget(false, false)
	if err != nil {
		return err
	}

	preFlavoringText := fmt.Sprintf("Setting label(s) for %s {{.ResourceName}} as {{.User}}...", strings.ToLower(cmd.RequiredArgs.ResourceType))
	cmd.UI.DisplayTextWithFlavor(
		preFlavoringText,
		map[string]interface{}{
			"ResourceName": cmd.RequiredArgs.ResourceName,
			"User":         username,
		},
	)

	warnings, err := cmd.Actor.UpdateStackLabelsByStackName(cmd.RequiredArgs.ResourceName, labels)
	cmd.UI.DisplayWarnings(warnings)

	return err
}

func (cmd SetLabelCommand) executeDomain(username string, labels map[string]types.NullString) error {
	preFlavoringText := fmt.Sprintf("Setting label(s) for %s {{.ResourceName}} as {{.User}}...", strings.ToLower(cmd.RequiredArgs.ResourceType))
	cmd.UI.DisplayTextWithFlavor(
		preFlavoringText,
		map[string]interface{}{
			"ResourceName": cmd.RequiredArgs.ResourceName,
			"User":         username,
		},
	)

	warnings, err := cmd.Actor.UpdateDomainLabelsByDomainName(cmd.RequiredArgs.ResourceName, labels)
	cmd.UI.DisplayWarnings(warnings)
	return err
}

func (cmd SetLabelCommand) executeServiceBroker(username string, labels map[string]types.NullString) error {
	if err := cmd.SharedActor.CheckTarget(false, false); err != nil {
		return err
	}

	preFlavoringText := fmt.Sprintf("Setting label(s) for %s {{.ResourceName}} as {{.User}}...", strings.ToLower(cmd.RequiredArgs.ResourceType))
	cmd.UI.DisplayTextWithFlavor(
		preFlavoringText,
		map[string]interface{}{
			"ResourceName": cmd.RequiredArgs.ResourceName,
			"User":         username,
		},
	)

	warnings, err := cmd.Actor.UpdateServiceBrokerLabelsByServiceBrokerName(cmd.RequiredArgs.ResourceName, labels)
	cmd.UI.DisplayWarnings(warnings)
	return err
}

func (cmd SetLabelCommand) canonicalResourceTypeForName() ResourceType {
	return ResourceType(strings.ToLower(cmd.RequiredArgs.ResourceType))
}

func (cmd SetLabelCommand) validateFlags() error {
	if cmd.BuildpackStack != "" && cmd.canonicalResourceTypeForName() != Buildpack {
		return translatableerror.ArgumentCombinationError{
			Args: []string{
				cmd.RequiredArgs.ResourceType, "--stack, -s",
			},
		}
	}
	return nil
}
<commit_msg>Reduce duplication in set-label command<commit_after>package v7

import (
	"fmt"
	"strings"

	"code.cloudfoundry.org/cli/actor/sharedaction"
	"code.cloudfoundry.org/cli/actor/v7action"
	"code.cloudfoundry.org/cli/command"
	"code.cloudfoundry.org/cli/command/flag"
	"code.cloudfoundry.org/cli/command/translatableerror"
	"code.cloudfoundry.org/cli/command/v7/shared"
	"code.cloudfoundry.org/cli/types"
	"code.cloudfoundry.org/clock"
)

//go:generate counterfeiter . SetLabelActor

type SetLabelActor interface {
	UpdateApplicationLabelsByApplicationName(string, string, map[string]types.NullString) (v7action.Warnings, error)
	UpdateBuildpackLabelsByBuildpackNameAndStack(string, string, map[string]types.NullString) (v7action.Warnings, error)
	UpdateDomainLabelsByDomainName(string, map[string]types.NullString) (v7action.Warnings, error)
	UpdateOrganizationLabelsByOrganizationName(string, map[string]types.NullString) (v7action.Warnings, error)
	UpdateRouteLabels(string, string, map[string]types.NullString) (v7action.Warnings, error)
	UpdateSpaceLabelsBySpaceName(string, string, map[string]types.NullString) (v7action.Warnings, error)
	UpdateStackLabelsByStackName(string, map[string]types.NullString) (v7action.Warnings, error)
	UpdateServiceBrokerLabelsByServiceBrokerName(string, map[string]types.NullString) (v7action.Warnings, error)
}

type SetLabelCommand struct {
	RequiredArgs    flag.SetLabelArgs `positional-args:"yes"`
	usage           interface{}       `usage:"CF_NAME set-label RESOURCE RESOURCE_NAME KEY=VALUE...\n\nEXAMPLES:\n   cf set-label app dora env=production\n   cf set-label org business pci=true public-facing=false\n   cf set-label buildpack go_buildpack go=1.12 -s cflinuxfs3\n\nRESOURCES:\n   app\n   buildpack\n   domain\n   org\n   route\n   service-broker\n   space\n   stack"`
	relatedCommands interface{}       `related_commands:"labels, unset-label"`
	BuildpackStack  string            `long:"stack" short:"s" description:"Specify stack to disambiguate buildpacks with the same name"`

	UI          command.UI
	Config      command.Config
	SharedActor command.SharedActor
	Actor       SetLabelActor

	username string
	labels   map[string]types.NullString
}

func (cmd *SetLabelCommand) Setup(config command.Config, ui command.UI) error {
	cmd.UI = ui
	cmd.Config = config
	cmd.SharedActor = sharedaction.NewActor(config)
	ccClient, _, err := shared.GetNewClientsAndConnectToCF(config, ui, "")
	if err != nil {
		return err
	}
	cmd.Actor = v7action.NewActor(ccClient, config, nil, nil, clock.NewClock())
	return nil
}

func (cmd SetLabelCommand) Execute(args []string) error {
	cmd.labels = make(map[string]types.NullString)
	for _, label := range cmd.RequiredArgs.Labels {
		parts := strings.SplitN(label, "=", 2)
		if len(parts) < 2 {
			return fmt.Errorf("Metadata error: no value provided for label '%s'", label)
		}
		cmd.labels[parts[0]] = types.NewNullString(parts[1])
	}

	var err error
	cmd.username, err = cmd.Config.CurrentUserName()
	if err != nil {
		return err
	}

	if err = cmd.validateFlags(); err != nil {
		return err
	}

	switch cmd.canonicalResourceTypeForName() {
	case App:
		err = cmd.executeApp()
	case Buildpack:
		err = cmd.executeBuildpack()
	case Domain:
		err = cmd.executeDomain()
	case Org:
		err = cmd.executeOrg()
	case Route:
		err = cmd.executeRoute()
	case Space:
		err = cmd.executeSpace()
	case Stack:
		err = cmd.executeStack()
	case ServiceBroker:
		err = cmd.executeServiceBroker()
	default:
		err = fmt.Errorf("Unsupported resource type of '%s'", cmd.RequiredArgs.ResourceType)
	}

	if err != nil {
		return err
	}

	cmd.UI.DisplayOK()
	return nil
}

func (cmd SetLabelCommand) executeApp() error {
	err := cmd.SharedActor.CheckTarget(true, true)
	if err != nil {
		return err
	}

	cmd.displayMessageWithOrgAndSpace()

	warnings, err := cmd.Actor.UpdateApplicationLabelsByApplicationName(
		cmd.RequiredArgs.ResourceName,
		cmd.Config.TargetedSpace().GUID,
		cmd.labels,
	)

	cmd.UI.DisplayWarnings(warnings)

	return err
}

func (cmd SetLabelCommand) executeBuildpack() error {
	err := cmd.SharedActor.CheckTarget(false, false)
	if err != nil {
		return err
	}

	preFlavoringText := fmt.Sprintf("Setting label(s) for %s {{.ResourceName}} as {{.User}}...", strings.ToLower(cmd.RequiredArgs.ResourceType))
	if cmd.BuildpackStack != "" {
		preFlavoringText = fmt.Sprintf("Setting label(s) for %s {{.ResourceName}} with stack {{.StackName}} as {{.User}}...", strings.ToLower(cmd.RequiredArgs.ResourceType))
	}

	cmd.UI.DisplayTextWithFlavor(
		preFlavoringText,
		map[string]interface{}{
			"ResourceName": cmd.RequiredArgs.ResourceName,
			"StackName":    cmd.BuildpackStack,
			"User":         cmd.username,
		},
	)

	warnings, err := cmd.Actor.UpdateBuildpackLabelsByBuildpackNameAndStack(cmd.RequiredArgs.ResourceName, cmd.BuildpackStack, cmd.labels)
	cmd.UI.DisplayWarnings(warnings)

	return err
}

func (cmd SetLabelCommand) executeOrg() error {
	err := cmd.SharedActor.CheckTarget(false, false)
	if err != nil {
		return err
	}

	cmd.displayMessage()

	warnings, err := cmd.Actor.UpdateOrganizationLabelsByOrganizationName(
		cmd.RequiredArgs.ResourceName,
		cmd.labels,
	)
	cmd.UI.DisplayWarnings(warnings)

	return err
}

func (cmd SetLabelCommand) executeRoute() error {
	err := cmd.SharedActor.CheckTarget(true, true)
	if err != nil {
		return err
	}

	cmd.displayMessageWithOrgAndSpace()

	warnings, err := cmd.Actor.UpdateRouteLabels(
		cmd.RequiredArgs.ResourceName,
		cmd.Config.TargetedSpace().GUID,
		cmd.labels,
	)
	cmd.UI.DisplayWarnings(warnings)

	return err
}

func (cmd SetLabelCommand) executeSpace() error {
	err := cmd.SharedActor.CheckTarget(true, false)
	if err != nil {
		return err
	}

	spaceName := cmd.RequiredArgs.ResourceName

	preFlavoringText := fmt.Sprintf("Setting label(s) for %s {{.ResourceName}} in org {{.OrgName}} as {{.User}}...", strings.ToLower(cmd.RequiredArgs.ResourceType))
	cmd.UI.DisplayTextWithFlavor(
		preFlavoringText,
		map[string]interface{}{
			"ResourceName": spaceName,
			"OrgName":      cmd.Config.TargetedOrganization().Name,
			"User":         cmd.username,
		},
	)

	warnings, err := cmd.Actor.UpdateSpaceLabelsBySpaceName(spaceName,
		cmd.Config.TargetedOrganization().GUID,
		cmd.labels)

	cmd.UI.DisplayWarnings(warnings)

	return err
}

func (cmd SetLabelCommand) executeStack() error {
	err := cmd.SharedActor.CheckTarget(false, false)
	if err != nil {
		return err
	}

	cmd.displayMessage()

	warnings, err := cmd.Actor.UpdateStackLabelsByStackName(cmd.RequiredArgs.ResourceName, cmd.labels)
	cmd.UI.DisplayWarnings(warnings)

	return err
}

func (cmd SetLabelCommand) executeDomain() error {
	cmd.displayMessage()

	warnings, err := cmd.Actor.UpdateDomainLabelsByDomainName(cmd.RequiredArgs.ResourceName, cmd.labels)
	cmd.UI.DisplayWarnings(warnings)
	return err
}

func (cmd SetLabelCommand) executeServiceBroker() error {
	if err := cmd.SharedActor.CheckTarget(false, false); err != nil {
		return err
	}

	cmd.displayMessage()

	warnings, err := cmd.Actor.UpdateServiceBrokerLabelsByServiceBrokerName(cmd.RequiredArgs.ResourceName, cmd.labels)
	cmd.UI.DisplayWarnings(warnings)
	return err
}

func (cmd SetLabelCommand) canonicalResourceTypeForName() ResourceType {
	return ResourceType(strings.ToLower(cmd.RequiredArgs.ResourceType))
}

func (cmd SetLabelCommand) validateFlags() error {
	if cmd.BuildpackStack != "" && cmd.canonicalResourceTypeForName() != Buildpack {
		return translatableerror.ArgumentCombinationError{
			Args: []string{
				cmd.RequiredArgs.ResourceType, "--stack, -s",
			},
		}
	}
	return nil
}

func (cmd SetLabelCommand) displayMessage() {
	preFlavoringText := fmt.Sprintf("Setting label(s) for %s {{.ResourceName}} as {{.User}}...", strings.ToLower(cmd.RequiredArgs.ResourceType))
	cmd.UI.DisplayTextWithFlavor(
		preFlavoringText,
		map[string]interface{}{
			"ResourceName": cmd.RequiredArgs.ResourceName,
			"User":         cmd.username,
		},
	)
}

func (cmd SetLabelCommand) displayMessageWithOrgAndSpace() {
	preFlavoringText := fmt.Sprintf("Setting label(s) for %s {{.ResourceName}} in org {{.OrgName}} / space {{.SpaceName}} as {{.User}}...", strings.ToLower(cmd.RequiredArgs.ResourceType))
	cmd.UI.DisplayTextWithFlavor(
		preFlavoringText,
		map[string]interface{}{
			"ResourceName": cmd.RequiredArgs.ResourceName,
			"OrgName":      cmd.Config.TargetedOrganization().Name,
			"SpaceName":    cmd.Config.TargetedSpace().Name,
			"User":         cmd.username,
		},
	)
}
<|endoftext|>


<commit_before>package test

import (
	"bytes"
	"encoding/base64"
	"encoding/json"
	"fmt"
	"math/big"
	"reflect"
	"runtime"
	"strings"
	"testing"
)

func fatalf(t *testing.T, format string, args ...interface{}) {
	fmt.Printf("\t"+format+"\n", args...)
	t.FailNow()
}

// Return short format caller info for printing errors, so errors don't all
// appear to come from test-tools.go.
func caller() string {
	_, file, line, _ := runtime.Caller(2)
	splits := strings.Split(file, "/")
	filename := splits[len(splits)-1]
	return fmt.Sprintf("%s:%d:", filename, line)
}

// Assert a boolean
func Assert(t *testing.T, result bool, message string) {
	if !result {
		fatalf(t, "%s %s", caller(), message)
	}
}

// AssertNotNil checks an object to be non-nil
func AssertNotNil(t *testing.T, obj interface{}, message string) {
	if obj == nil {
		fatalf(t, "%s %s", caller(), message)
	}
}

// AssertNotError checks that err is nil
func AssertNotError(t *testing.T, err error, message string) {
	if err != nil {
		fatalf(t, "%s %s: %s", caller(), message, err)
	}
}

// AssertError checks that err is non-nil
func AssertError(t *testing.T, err error, message string) {
	if err == nil {
		fatalf(t, "%s %s: expected error but received none", caller(), message)
	}
}

// AssertEquals uses the equality operator (==) to measure one and two
func AssertEquals(t *testing.T, one interface{}, two interface{}) {
	if one != two {
		fatalf(t, "%s [%v] != [%v]", caller(), one, two)
	}
}

// AssertDeepEquals uses the reflect.DeepEqual method to measure one and two
func AssertDeepEquals(t *testing.T, one interface{}, two interface{}) {
	if !reflect.DeepEqual(one, two) {
		fatalf(t, "%s [%+v] !(deep)= [%+v]", caller(), one, two)
	}
}

// AssertMarshaledEquals marshals one and two to JSON, and then uses
// the equality operator to measure them
func AssertMarshaledEquals(t *testing.T, one interface{}, two interface{}) {
	oneJSON, err := json.Marshal(one)
	AssertNotError(t, err, "Could not marshal 1st argument")
	twoJSON, err := json.Marshal(two)
	AssertNotError(t, err, "Could not marshal 2nd argument")

	if !bytes.Equal(oneJSON, twoJSON) {
		fatalf(t, "%s [%s] !(json)= [%s]", caller(), oneJSON, twoJSON)
	}
}

// AssertNotEquals uses the equality operator to measure that one and two
// are different
func AssertNotEquals(t *testing.T, one interface{}, two interface{}) {
	if one == two {
		fatalf(t, "%s [%v] == [%v]", caller(), one, two)
	}
}

// AssertByteEquals uses bytes.Equal to measure one and two for equality.
func AssertByteEquals(t *testing.T, one []byte, two []byte) {
	if !bytes.Equal(one, two) {
		fatalf(t, "%s Byte [%s] != [%s]",
			caller(),
			base64.StdEncoding.EncodeToString(one),
			base64.StdEncoding.EncodeToString(two))
	}
}

// AssertIntEquals uses the equality operator to measure one and two.
func AssertIntEquals(t *testing.T, one int, two int) {
	if one != two {
		fatalf(t, "%s Int [%d] != [%d]", caller(), one, two)
	}
}

// AssertBigIntEquals uses the big.Int.cmp() method to measure whether
// one and two are equal
func AssertBigIntEquals(t *testing.T, one *big.Int, two *big.Int) {
	if one.Cmp(two) != 0 {
		fatalf(t, "%s Int [%d] != [%d]", caller(), one, two)
	}
}

// AssertContains determines whether needle can be found in haystack
func AssertContains(t *testing.T, haystack string, needle string) {
	if !strings.Contains(haystack, needle) {
		fatalf(t, "%s String [%s] does not contain [%s]", caller(), haystack, needle)
	}
}

// AssertNotContains determines if needle is not found in haystack
func AssertNotContains(t *testing.T, haystack string, needle string) {
	if strings.Contains(haystack, needle) {
		fatalf(t, "%s String [%s] contains [%s]", caller(), haystack, needle)
	}
}

// AssertSeverity determines if a string matches the Severity formatting
func AssertSeverity(t *testing.T, data string, severity int) {
	expected := fmt.Sprintf("\"severity\":%d", severity)
	AssertContains(t, data, expected)
}

// AssertBetween determines if a is between b and c
func AssertBetween(t *testing.T, a, b, c int64) {
	if a < b || a > c {
		fatalf(t, "%d is not between %d and %d", a, b, c)
	}
}
<commit_msg>Change `%v` to `%#v` in AssertEquals error msg. (#1983)<commit_after>package test

import (
	"bytes"
	"encoding/base64"
	"encoding/json"
	"fmt"
	"math/big"
	"reflect"
	"runtime"
	"strings"
	"testing"
)

func fatalf(t *testing.T, format string, args ...interface{}) {
	fmt.Printf("\t"+format+"\n", args...)
	t.FailNow()
}

// Return short format caller info for printing errors, so errors don't all
// appear to come from test-tools.go.
func caller() string {
	_, file, line, _ := runtime.Caller(2)
	splits := strings.Split(file, "/")
	filename := splits[len(splits)-1]
	return fmt.Sprintf("%s:%d:", filename, line)
}

// Assert a boolean
func Assert(t *testing.T, result bool, message string) {
	if !result {
		fatalf(t, "%s %s", caller(), message)
	}
}

// AssertNotNil checks an object to be non-nil
func AssertNotNil(t *testing.T, obj interface{}, message string) {
	if obj == nil {
		fatalf(t, "%s %s", caller(), message)
	}
}

// AssertNotError checks that err is nil
func AssertNotError(t *testing.T, err error, message string) {
	if err != nil {
		fatalf(t, "%s %s: %s", caller(), message, err)
	}
}

// AssertError checks that err is non-nil
func AssertError(t *testing.T, err error, message string) {
	if err == nil {
		fatalf(t, "%s %s: expected error but received none", caller(), message)
	}
}

// AssertEquals uses the equality operator (==) to measure one and two
func AssertEquals(t *testing.T, one interface{}, two interface{}) {
	if one != two {
		fatalf(t, "%s %#v != %#v", caller(), one, two)
	}
}

// AssertDeepEquals uses the reflect.DeepEqual method to measure one and two
func AssertDeepEquals(t *testing.T, one interface{}, two interface{}) {
	if !reflect.DeepEqual(one, two) {
		fatalf(t, "%s [%+v] !(deep)= [%+v]", caller(), one, two)
	}
}

// AssertMarshaledEquals marshals one and two to JSON, and then uses
// the equality operator to measure them
func AssertMarshaledEquals(t *testing.T, one interface{}, two interface{}) {
	oneJSON, err := json.Marshal(one)
	AssertNotError(t, err, "Could not marshal 1st argument")
	twoJSON, err := json.Marshal(two)
	AssertNotError(t, err, "Could not marshal 2nd argument")

	if !bytes.Equal(oneJSON, twoJSON) {
		fatalf(t, "%s [%s] !(json)= [%s]", caller(), oneJSON, twoJSON)
	}
}

// AssertNotEquals uses the equality operator to measure that one and two
// are different
func AssertNotEquals(t *testing.T, one interface{}, two interface{}) {
	if one == two {
		fatalf(t, "%s %#v == %#v", caller(), one, two)
	}
}

// AssertByteEquals uses bytes.Equal to measure one and two for equality.
func AssertByteEquals(t *testing.T, one []byte, two []byte) {
	if !bytes.Equal(one, two) {
		fatalf(t, "%s Byte [%s] != [%s]",
			caller(),
			base64.StdEncoding.EncodeToString(one),
			base64.StdEncoding.EncodeToString(two))
	}
}

// AssertIntEquals uses the equality operator to measure one and two.
func AssertIntEquals(t *testing.T, one int, two int) {
	if one != two {
		fatalf(t, "%s Int [%d] != [%d]", caller(), one, two)
	}
}

// AssertBigIntEquals uses the big.Int.cmp() method to measure whether
// one and two are equal
func AssertBigIntEquals(t *testing.T, one *big.Int, two *big.Int) {
	if one.Cmp(two) != 0 {
		fatalf(t, "%s Int [%d] != [%d]", caller(), one, two)
	}
}

// AssertContains determines whether needle can be found in haystack
func AssertContains(t *testing.T, haystack string, needle string) {
	if !strings.Contains(haystack, needle) {
		fatalf(t, "%s String [%s] does not contain [%s]", caller(), haystack, needle)
	}
}

// AssertNotContains determines if needle is not found in haystack
func AssertNotContains(t *testing.T, haystack string, needle string) {
	if strings.Contains(haystack, needle) {
		fatalf(t, "%s String [%s] contains [%s]", caller(), haystack, needle)
	}
}

// AssertSeverity determines if a string matches the Severity formatting
func AssertSeverity(t *testing.T, data string, severity int) {
	expected := fmt.Sprintf("\"severity\":%d", severity)
	AssertContains(t, data, expected)
}

// AssertBetween determines if a is between b and c
func AssertBetween(t *testing.T, a, b, c int64) {
	if a < b || a > c {
		fatalf(t, "%d is not between %d and %d", a, b, c)
	}
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"github.com/miekg/dns"
	"gloon/mem_rs"
	"gloon/record_set"
	"gloon/redis_rs"
	"log"
	"regexp"
)

type Server struct {
	*dns.Server
	*record_set.RecordSet
	resolver *Resolver
	settings *Settings
}

var split_rex = regexp.MustCompile("[:= ]")

func NewServer(addr string, settings *Settings) (s *Server, err error) {
	s = &Server{}
	s.settings = settings
	s.Server = &dns.Server{Addr: addr, Net: "udp"}
	dns.HandleFunc(".", func(w dns.ResponseWriter, r *dns.Msg) {
		s.handleDnsRequest(w, r)
	})
	s.resolver, err = NewResolver(settings)
	for _, v := range settings.Hostnames {
		parts := split_rex.Split(v, -1)
		if len(parts) == 2 {
			s.Records.Put(dns.TypeA, parts[0], parts[1])
		}
	}
	switch settings.Store {
	case "redis":
		server.record_set, err = redis_rs.Create("127.0.0.1:6379", redis_rs.RedisRecordStoreOpts{})
		if err != nil {
			log.Fatalf("Unable to create redis record set: %s", err.Error())
		}
	default:
		server.record_set = mem_rs.Create()
	}
	return
}

func (s *Server) handleDnsRequest(w dns.ResponseWriter, r *dns.Msg) {
	defer func() {
		if r := recover(); r != nil {
			handlePanic(r)
		}
	}()
	m := new(dns.Msg)
	m.SetReply(r)
	m.Compress = false
	switch r.Opcode {
	case dns.OpcodeQuery:
		if s.processQuery(m) {
			w.WriteMsg(m)
			return
		}
	}
	if s.settings.DisableForwarding {
		log.Printf("WARNING: name not found and forwarding disabled")
		m.Rcode = dns.RcodeNameError
		w.WriteMsg(m)
		return
	}
	resp, err := s.resolver.Lookup(r)
	if err == nil {
		w.WriteMsg(resp)
	} else {
		log.Printf("Resolver err: %s", err.Error())
	}
}

func (s *Server) processQuery(m *dns.Msg) bool {
	answers := 0

	for _, q := range m.Question {
		var rr dns.RR
		switch q.Qtype {
		case dns.TypeA:
			ip := s.Get(dns.TypeA, q.Name)
			if ip != "" {
				rr, _ = dns.NewRR(fmt.Sprintf("%s A %s", q.Name, ip))
			}
		case dns.TypePTR:
			host := s.Get(dns.TypePTR, q.Name)
			if host != "" {
				rr, _ = dns.NewRR(fmt.Sprintf("%s PTR %s", q.Name, host))
			}
		case dns.TypeAAAA: // For now, if we have an A record, for a name, always NXDOMAIN for corresponding AAAA records
			ip := s.Get(dns.TypeAAAA, q.Name)
			if ip != "" {
				rr, _ = dns.NewRR(fmt.Sprintf("%s AAAA %s", q.Name, ip))
			}
		default:
			return false // If we get a question we can't answer, bail
		}
		if rr != nil {
			m.Answer = append(m.Answer, rr)
			answers++
		}
	}
	if answers > 0 {
		return true
	}
	return false
}
<commit_msg>set up store correctly<commit_after>package main

import (
	"fmt"
	"github.com/miekg/dns"
	"gloon/mem_rs"
	"gloon/record_set"
	"gloon/redis_rs"
	"log"
	"regexp"
)

type Server struct {
	*dns.Server
	*record_set.RecordSet
	resolver *Resolver
	settings *Settings
}

var split_rex = regexp.MustCompile("[:= ]")

func NewServer(addr string, settings *Settings) (s *Server, err error) {
	s = &Server{}
	s.settings = settings

	// Set up the record store
	var store record_set.RecordStore
	switch settings.Store {
	case "redis":
		store, err = redis_rs.Create("127.0.0.1:6379", &redis_rs.RedisRecordStoreOpts{})
		if err != nil {
			log.Fatalf("Unable to create redis record set: %s", err.Error())
		}
	default:
		store = mem_rs.Create()
	}
	s.RecordSet = record_set.Create(store)
	s.Server = &dns.Server{Addr: addr, Net: "udp"}
	dns.HandleFunc(".", func(w dns.ResponseWriter, r *dns.Msg) {
		s.handleDnsRequest(w, r)
	})
	s.resolver, err = NewResolver(settings)
	for _, v := range settings.Hostnames {
		parts := split_rex.Split(v, -1)
		if len(parts) == 2 {
			s.RecordSet.Put(dns.TypeA, parts[0], parts[1])
		}
	}
	return
}

func (s *Server) handleDnsRequest(w dns.ResponseWriter, r *dns.Msg) {
	defer func() {
		if r := recover(); r != nil {
			handlePanic(r)
		}
	}()
	m := new(dns.Msg)
	m.SetReply(r)
	m.Compress = false
	switch r.Opcode {
	case dns.OpcodeQuery:
		if s.processQuery(m) {
			w.WriteMsg(m)
			return
		}
	}
	if s.settings.DisableForwarding {
		log.Printf("WARNING: name not found and forwarding disabled")
		m.Rcode = dns.RcodeNameError
		w.WriteMsg(m)
		return
	}
	resp, err := s.resolver.Lookup(r)
	if err == nil {
		w.WriteMsg(resp)
	} else {
		log.Printf("Resolver err: %s", err.Error())
	}
}

func (s *Server) processQuery(m *dns.Msg) bool {
	answers := 0

	for _, q := range m.Question {
		var rr dns.RR
		switch q.Qtype {
		case dns.TypeA:
			ip := s.Get(dns.TypeA, q.Name)
			if ip != "" {
				rr, _ = dns.NewRR(fmt.Sprintf("%s A %s", q.Name, ip))
			}
		case dns.TypePTR:
			host := s.Get(dns.TypePTR, q.Name)
			if host != "" {
				rr, _ = dns.NewRR(fmt.Sprintf("%s PTR %s", q.Name, host))
			}
		case dns.TypeAAAA: // For now, if we have an A record, for a name, always NXDOMAIN for corresponding AAAA records
			ip := s.Get(dns.TypeAAAA, q.Name)
			if ip != "" {
				rr, _ = dns.NewRR(fmt.Sprintf("%s AAAA %s", q.Name, ip))
			}
		default:
			return false // If we get a question we can't answer, bail
		}
		if rr != nil {
			m.Answer = append(m.Answer, rr)
			answers++
		}
	}
	if answers > 0 {
		return true
	}
	return false
}
<|endoftext|>


<commit_before>package httpfstream

import (
	"bytes"
	"io"
	"net/http"
	"strings"
	"testing"
)

type uploadTest struct {
	path string
	data io.Reader

	// error responses
	statusCode int
	msg        string
}

func TestUpload(t *testing.T) {
	tests := []uploadTest{
		{path: "/foo", data: bytes.NewReader([]byte("bar")), statusCode: http.StatusOK},
		{path: "/foo", statusCode: http.StatusOK},

		{path: "/", statusCode: http.StatusBadRequest, msg: "path must not end with '/'"},
		{path: "/..", statusCode: http.StatusMovedPermanently},
		{path: "/../foo", statusCode: http.StatusMovedPermanently},
	}
	for _, test := range tests {
		testUpload(t, test)
	}
}

func testUpload(t *testing.T, test uploadTest) {
	label := test.path

	server := newTestServer()
	defer server.close()

	req, err := http.NewRequest("PUT", server.URL+test.path, test.data)
	if err != nil {
		t.Fatalf("%s: NewRequest: %s", label, err)
	}

	resp, err := http.DefaultClient.Do(req)
	if err != nil {
		t.Fatalf("%s: Do: %s", label, err)
	}

	if test.statusCode != resp.StatusCode {
		t.Errorf("%s: want StatusCode == %d, got %d", label, test.statusCode, resp.StatusCode)
	}

	if test.statusCode >= 200 && test.statusCode <= 299 {
		_, err = server.fs.Stat(test.path)
		if err != nil {
			t.Errorf("%s: Stat: %s", label, err)
		}
	} else {
		msg := strings.TrimSpace(string(readAll(t, resp.Body)))
		if test.msg != msg {
			t.Errorf("%s: want error message %q, got %q", label, test.msg, msg)
		}
	}
}
<commit_msg>TestUpload bytes.Buffer data source<commit_after>package httpfstream

import (
	"bytes"
	"io"
	"net/http"
	"strings"
	"testing"
)

type uploadTest struct {
	path string
	data io.Reader

	// error responses
	statusCode int
	msg        string
}

func TestUpload(t *testing.T) {
	tests := []uploadTest{
		{path: "/foo", data: bytes.NewReader([]byte("bar")), statusCode: http.StatusOK},
		{path: "/foo", data: bytes.NewBuffer([]byte("bar")), statusCode: http.StatusOK},
		{path: "/foo", statusCode: http.StatusOK},

		{path: "/", statusCode: http.StatusBadRequest, msg: "path must not end with '/'"},
		{path: "/..", statusCode: http.StatusMovedPermanently},
		{path: "/../foo", statusCode: http.StatusMovedPermanently},
	}
	for _, test := range tests {
		testUpload(t, test)
	}
}

func testUpload(t *testing.T, test uploadTest) {
	label := test.path

	server := newTestServer()
	defer server.close()

	req, err := http.NewRequest("PUT", server.URL+test.path, test.data)
	if err != nil {
		t.Fatalf("%s: NewRequest: %s", label, err)
	}

	resp, err := http.DefaultClient.Do(req)
	if err != nil {
		t.Fatalf("%s: Do: %s", label, err)
	}

	if test.statusCode != resp.StatusCode {
		t.Errorf("%s: want StatusCode == %d, got %d", label, test.statusCode, resp.StatusCode)
	}

	if test.statusCode >= 200 && test.statusCode <= 299 {
		_, err = server.fs.Stat(test.path)
		if err != nil {
			t.Errorf("%s: Stat: %s", label, err)
		}
	} else {
		msg := strings.TrimSpace(string(readAll(t, resp.Body)))
		if test.msg != msg {
			t.Errorf("%s: want error message %q, got %q", label, test.msg, msg)
		}
	}
}
<|endoftext|>


<commit_before>package main

import (
	"crypto/hmac"
	"crypto/sha1"
	"encoding/hex"
	"encoding/json"
	"errors"
	"flag"
	"io"
	"io/ioutil"
	"log"
	"net"
	"net/http"
	"net/url"
	"regexp"
	"runtime"
	"strings"
	"time"
)

var validReqHeaders = map[string]bool{
	"Accept":            true,
	"Accept-Charset":    true,
	"Accept-Encoding":   true,
	"Cache-Control":     true,
	"If-Modified_Since": true,
}

func noRedirect(req *http.Request, via []*http.Request) error {
	return errors.New("Redirect")
}

func validateURL(path string, key []byte) (surl string, valid bool) {
	pathParts := strings.SplitN(path[1:], "/", 3)
	valid = false
	if len(pathParts) != 2 {
		log.Println("Bad path format", pathParts)
		return
	}
	hexdig, hexurl := pathParts[0], pathParts[1]
	urlBytes, err := hex.DecodeString(hexurl)
	if err != nil {
		log.Println("Bad Hex Decode", hexurl)
		return
	}
	surl = string(urlBytes)
	//log.Println(surl)
	mac := hmac.New(sha1.New, key)
	mac.Write([]byte(surl))
	macSum := hex.EncodeToString(mac.Sum([]byte{}))
	if macSum != hexdig {
		log.Printf("Bad signature: %s != %s", macSum, hexdig)
		return
	}
	valid = true
	return
}

type ProxyHandler struct {
	Transport       *http.Transport
	HMacKey         []byte
	RegexpAllowlist []*regexp.Regexp
	RegexpDenylist  []*regexp.Regexp
}

func (p *ProxyHandler) ServeHTTP(w http.ResponseWriter, req *http.Request) {
	//log.Println("Request:", req.URL)

	// do fiddly things
	if req.Method != "GET" {
		log.Println("Something other than GET received", req.Method)
		http.Error(w, "Method Not Implemented", http.StatusNotImplemented)
		return
	}

	surl, ok := validateURL(req.URL.Path, p.HMacKey)
	if !ok {
		http.Error(w, "Bad Signature", http.StatusForbidden)
		return
	}
	//log.Println("URL:", surl)

	u, err := url.Parse(surl)
	if err != nil {
		log.Println(err)
		http.Error(w, "Bad url", http.StatusBadRequest)
		return
	}

	if u.Host == "" {
		http.Error(w, "Bad url", http.StatusNotFound)
		return
	}

	// if Allowlist is set, require match
	matchFound := true
	if len(p.RegexpAllowlist) > 0 {
		matchFound = false
		for _, rgx := range p.RegexpAllowlist {
			if rgx.MatchString(u.Host) {
				matchFound = true
			}
		}
	}
	if !matchFound {
		http.Error(w, "Allowlist host failure", http.StatusNotFound)
		return
	}

	// filter out Denylist urls based on regexes. Do this second
	// as Denylist takes precedence
	for _, rgx := range p.RegexpDenylist {
		if rgx.MatchString(u.Host) {
			http.Error(w, "Denylist host failure", http.StatusNotFound)
			return
		}
	}

	nreq, err := http.NewRequest("GET", surl, nil)
	if err != nil {
		log.Println("Something weird happened")
		http.Error(w, "Error Fetching Resource", http.StatusNotFound)
		return
	}

	// filter headers
	for hdr, val := range req.Header {
		if validReqHeaders[hdr] {
			nreq.Header[hdr] = val
		}
	}
	nreq.Header.Add("connection", "close")
	nreq.Header.Add("user-agent", "pew pew pew")

	resp, err := p.Transport.RoundTrip(nreq)
	if err != nil {
		log.Println("Could not connect to endpoint", err)
		if strings.Contains(err.Error(), "timeout") {
			http.Error(w, "Error Fetching Resource", http.StatusBadGateway)
		} else {
			http.Error(w, "Error Fetching Resource", http.StatusNotFound)
		}
		return
	}
	defer resp.Body.Close()
	defer p.Transport.CloseIdleConnections()

	// check for too large a response
	if resp.ContentLength > 5242880 {
		log.Println("Content length exceeded")
		http.Error(w, "Content length exceeded", http.StatusBadRequest)
		return
	}

	// check for redirects. we do not follow.
	if resp.StatusCode >= 300 && resp.StatusCode <= 307 {
		log.Println("Refusing to follow redirects")
		http.Error(w, "Refusing to follow redirects", http.StatusNotFound)
		return
	}

	// check content type
	ct, ok := resp.Header[http.CanonicalHeaderKey("content-type")]
	if !ok || ct[0][:6] != "image/" {
		log.Println("Non-Image content-type returned")
		http.Error(w, "Non-Image content-type returned", http.StatusBadRequest)
		return
	}

	for hdr, val := range resp.Header {
		h := w.Header()
		h[hdr] = val
	}
	h := w.Header()
	h.Add("X-Content-Type-Options", "nosniff")

	io.Copy(w, resp.Body)
	//log.Println(req, resp.StatusCode)
}

var hmacKeyFlag = flag.String("hmacKey", "", "HMAC Key")
var configFileFlag = flag.String("configFile", "", "JSON Config File")
var bindAddress = flag.String("bindAddress", "0.0.0.0:8080", "Address:Port to bind to")

type configParams struct {
	HmacKey   string
	Allowlist []string
	Denylist  []string
}

func main() {
	runtime.GOMAXPROCS(runtime.NumCPU())
	flag.Parse()

	config := &configParams{}
	if *configFileFlag != "" {
		b, err := ioutil.ReadFile(*configFileFlag)
		if err != nil {
			log.Fatal("Could not read configFile", err)
		}
		err = json.Unmarshal(b, &config)
		if err != nil {
			log.Fatal("Could not parse configFile", err)
		}
	}

	// flags override config file
	if *hmacKeyFlag != "" {
		config.HmacKey = *hmacKeyFlag
	}
	//log.Println("Config", config.HmacKey, config.Allowlist, config.Denylist)

	tr := &http.Transport{
		Dial: func(netw, addr string) (net.Conn, error) {
			// 2 second timeout on requests
			timeout := time.Second * time.Duration(2)
			c, err := net.DialTimeout(netw, addr, timeout)
			if err != nil {
				return nil, err
			}
			// also set time limit on reading
			c.SetDeadline(time.Now().Add(timeout))
			return c, nil
		}}
	proxy := &ProxyHandler{
		Transport: tr,
		HMacKey:   []byte(config.HmacKey)}

	// build/compile regex
	proxy.RegexpAllowlist = make([]*regexp.Regexp, 0)
	proxy.RegexpDenylist = make([]*regexp.Regexp, 0)

	var c *regexp.Regexp
	var err error
	for _, v := range config.Denylist {
		c, err = regexp.Compile(v)
		if err != nil {
			log.Fatal(err)
		}
		proxy.RegexpDenylist = append(proxy.RegexpDenylist, c)
	}
	for _, v := range config.Allowlist {
		c, err = regexp.Compile(v)
		if err != nil {
			log.Fatal(err)
		}
		proxy.RegexpAllowlist = append(proxy.RegexpAllowlist, c)
	}

	http.Handle("/favicon.ico", http.NotFoundHandler())
	http.Handle("/", proxy)
	log.Println("Starting server on", *bindAddress)
	err = http.ListenAndServe(*bindAddress, nil)
	if err != nil {
		log.Fatal("ListenAndServe: ", err.Error())
	}
}
<commit_msg>return original resp code, and reap transport wait conns in a goroutine<commit_after>package main

import (
	"crypto/hmac"
	"crypto/sha1"
	"encoding/hex"
	"encoding/json"
	"errors"
	"flag"
	"io"
	"io/ioutil"
	"log"
	"net"
	"net/http"
	"net/url"
	"regexp"
	"runtime"
	"strings"
	"time"
)

var validReqHeaders = map[string]bool{
	"Accept":            true,
	"Accept-Charset":    true,
	"Accept-Encoding":   true,
	"Accept-Language":   true,
	"Cache-Control":     true,
	"If-None-Match":     true,
	"If-Modified-Since": true,
}

func noRedirect(req *http.Request, via []*http.Request) error {
	return errors.New("Redirect")
}

func validateURL(path string, key []byte) (surl string, valid bool) {
	pathParts := strings.SplitN(path[1:], "/", 3)
	valid = false
	if len(pathParts) != 2 {
		log.Println("Bad path format", pathParts)
		return
	}
	hexdig, hexurl := pathParts[0], pathParts[1]
	urlBytes, err := hex.DecodeString(hexurl)
	if err != nil {
		log.Println("Bad Hex Decode", hexurl)
		return
	}
	surl = string(urlBytes)
	//log.Println(surl)
	mac := hmac.New(sha1.New, key)
	mac.Write([]byte(surl))
	macSum := hex.EncodeToString(mac.Sum([]byte{}))
	if macSum != hexdig {
		log.Printf("Bad signature: %s != %s", macSum, hexdig)
		return
	}
	valid = true
	return
}

type ProxyHandler struct {
	Transport       *http.Transport
	HMacKey         []byte
	RegexpAllowlist []*regexp.Regexp
	RegexpDenylist  []*regexp.Regexp
}

func (p *ProxyHandler) ServeHTTP(w http.ResponseWriter, req *http.Request) {
	//log.Println("Request:", req.URL)

	// do fiddly things
	if req.Method != "GET" {
		log.Println("Something other than GET received", req.Method)
		http.Error(w, "Method Not Implemented", http.StatusNotImplemented)
		return
	}

	surl, ok := validateURL(req.URL.Path, p.HMacKey)
	if !ok {
		http.Error(w, "Bad Signature", http.StatusForbidden)
		return
	}
	//log.Println("URL:", surl)

	u, err := url.Parse(surl)
	if err != nil {
		log.Println(err)
		http.Error(w, "Bad url", http.StatusBadRequest)
		return
	}

	if u.Host == "" {
		http.Error(w, "Bad url", http.StatusNotFound)
		return
	}

	// if Allowlist is set, require match
	matchFound := true
	if len(p.RegexpAllowlist) > 0 {
		matchFound = false
		for _, rgx := range p.RegexpAllowlist {
			if rgx.MatchString(u.Host) {
				matchFound = true
			}
		}
	}
	if !matchFound {
		http.Error(w, "Allowlist host failure", http.StatusNotFound)
		return
	}

	// filter out Denylist urls based on regexes. Do this second
	// as Denylist takes precedence
	for _, rgx := range p.RegexpDenylist {
		if rgx.MatchString(u.Host) {
			http.Error(w, "Denylist host failure", http.StatusNotFound)
			return
		}
	}

	nreq, err := http.NewRequest("GET", surl, nil)
	if err != nil {
		log.Println("Something weird happened")
		http.Error(w, "Error Fetching Resource", http.StatusNotFound)
		return
	}

	// filter headers
	for hdr, val := range req.Header {
		if validReqHeaders[hdr] {
			nreq.Header[hdr] = val
		}
	}
	nreq.Header.Add("connection", "close")
	nreq.Header.Add("user-agent", "pew pew pew")

	resp, err := p.Transport.RoundTrip(nreq)
	if err != nil {
		log.Println("Could not connect to endpoint", err)
		if strings.Contains(err.Error(), "timeout") {
			http.Error(w, "Error Fetching Resource", http.StatusBadGateway)
		} else {
			http.Error(w, "Error Fetching Resource", http.StatusNotFound)
		}
		return
	}
	defer resp.Body.Close()

	// check for too large a response
	if resp.ContentLength > 5242880 {
		log.Println("Content length exceeded")
		http.Error(w, "Content length exceeded", http.StatusBadRequest)
		return
	}

	// check for redirects. we do not follow.
	if resp.StatusCode >= 300 && resp.StatusCode <= 303 {
		log.Println("Refusing to follow redirects")
		http.Error(w, "Refusing to follow redirects", http.StatusNotFound)
		return
	}

	// check content type
	ct, ok := resp.Header[http.CanonicalHeaderKey("content-type")]
	// if not a 200, then we will just pass it along
	if resp.StatusCode == 200 {
		if !ok || ct[0][:6] != "image/" {
			log.Println("Non-Image content-type returned", u)
			http.Error(w, "Non-Image content-type returned",
					   http.StatusBadRequest)
			return
		}
	}

	for hdr, val := range resp.Header {
		h := w.Header()
		h[hdr] = val
	}
	h := w.Header()
	h.Add("X-Content-Type-Options", "nosniff")

	w.WriteHeader(resp.StatusCode)

	io.Copy(w, resp.Body)
	//log.Println(req, resp.StatusCode)
}

var hmacKeyFlag = flag.String("hmacKey", "", "HMAC Key")
var configFileFlag = flag.String("configFile", "", "JSON Config File")
var bindAddress = flag.String("bindAddress", "0.0.0.0:8080", "Address:Port to bind to")

type configParams struct {
	HmacKey   string
	Allowlist []string
	Denylist  []string
}

func main() {
	runtime.GOMAXPROCS(runtime.NumCPU())
	flag.Parse()

	config := &configParams{}
	if *configFileFlag != "" {
		b, err := ioutil.ReadFile(*configFileFlag)
		if err != nil {
			log.Fatal("Could not read configFile", err)
		}
		err = json.Unmarshal(b, &config)
		if err != nil {
			log.Fatal("Could not parse configFile", err)
		}
	}

	// flags override config file
	if *hmacKeyFlag != "" {
		config.HmacKey = *hmacKeyFlag
	}
	//log.Println("Config", config.HmacKey, config.Allowlist, config.Denylist)

	tr := &http.Transport{
		Dial: func(netw, addr string) (net.Conn, error) {
			// 2 second timeout on requests
			timeout := time.Second * 2
			c, err := net.DialTimeout(netw, addr, timeout)
			if err != nil {
				return nil, err
			}
			// also set time limit on reading
			c.SetDeadline(time.Now().Add(timeout))
			return c, nil
		}}

	// spawn an idle conn trimmer
	go func() {
		time.Sleep(5 * time.Minute)
		tr.CloseIdleConnections()
	}()

	proxy := &ProxyHandler{
		Transport: tr,
		HMacKey:   []byte(config.HmacKey)}

	// build/compile regex
	proxy.RegexpAllowlist = make([]*regexp.Regexp, 0)
	proxy.RegexpDenylist = make([]*regexp.Regexp, 0)

	var c *regexp.Regexp
	var err error
	for _, v := range config.Denylist {
		c, err = regexp.Compile(v)
		if err != nil {
			log.Fatal(err)
		}
		proxy.RegexpDenylist = append(proxy.RegexpDenylist, c)
	}
	for _, v := range config.Allowlist {
		c, err = regexp.Compile(v)
		if err != nil {
			log.Fatal(err)
		}
		proxy.RegexpAllowlist = append(proxy.RegexpAllowlist, c)
	}

	http.Handle("/favicon.ico", http.NotFoundHandler())
	http.Handle("/", proxy)
	log.Println("Starting server on", *bindAddress)
	err = http.ListenAndServe(*bindAddress, nil)
	if err != nil {
		log.Fatal("ListenAndServe: ", err.Error())
	}
}
<|endoftext|>