<commit_before>package main

import (
	"flag"
	"fmt"
	"github.com/rschoen/fireworks-server/lib"
	"log"
	"math/rand"
	"net/http"
	"regexp"
	"strconv"
	"strings"
	"time"
)

func (s *Server) handler(w http.ResponseWriter, r *http.Request) {
	// allow requests to come from anywhere, since clients can be wherever
	w.Header().Set("Access-Control-Allow-Origin", "*")

	// serve client HTTP responses, if it's turned on
	if len(r.URL.Path) < 5 || r.URL.Path[1:5] != "api/" {
		if s.httpServer == true {
			http.FileServer(http.Dir(s.clientDirectory)).ServeHTTP(w, r)
		}
		return
	}

	m, err := lib.DecodeMove(r.PostFormValue("data"))
	if err != "" {
		log.Println("Discarding malformed JSON message. Error: " + err)
		fmt.Fprintf(w, jsonError("Data sent was malformed."))
		return
	}

	var command = r.URL.Path[5:]
	var game *lib.Game
	for _, ongoingGame := range s.games {
		if ongoingGame.GameID == m.Game {
			game = ongoingGame
		}
	}

	// Authenticate user
	authResponse, authError := lib.Authenticate(m.Token)
	if authError != "" {
		log.Printf("Failed to authenticate player '%s' in game '%s'. Error: %s\n", m.Player, m.Game, authError)
		fmt.Fprintf(w, jsonError("Could not authorize user."))
		return
	}
	if authResponse.GetGoogleID() != m.Player {
		log.Printf("Authenticated player '%s' submitted move as player '%s' in game '%s'.", authResponse.GetGoogleID(), m.Player, m.Game, m.Game)
		fmt.Fprintf(w, jsonError("Authenticated as a different user."))
		return
	}

	if command == "list" {
		list := lib.GamesList{}
		for i, _ := range s.games {
			playerList := ""
			addGame := false
			for player, _ := range s.games[i].Players {
				playerList += s.games[i].Players[player].Name + ", "
				if s.games[i].Players[player].GoogleID == m.Player {
					addGame = true
				}
			}
			if playerList != "" {
				playerList = playerList[:len(playerList)-2]
			}
			game := lib.MinimalGame{Name: s.games[i].GameID, Players: playerList}

			if addGame {
				list.PlayersGames = append(list.PlayersGames, game)
			} else if s.games[i].State == lib.StateNotStarted && len(s.games[i].Players) < lib.MaxPlayers {
				list.OpenGames = append(list.OpenGames, game)
			}
		}

		encodedList, err := lib.EncodeList(list)
		if err != "" {
			log.Printf("Failed to encode game list. Error: %s\n", err)
			fmt.Fprintf(w, jsonError("Could not transmit game list to client."))
			return
		}
		fmt.Fprintf(w, encodedList)
		return
	}

	if command == "join" {
		// create game if it doesn't exist
		if game == nil {
			game = new(lib.Game)
			game.GameID = sanitizeAndTrim(m.Game, lib.MaxGameNameLength, false)
			var initializationError = game.Initialize()
			if initializationError != "" {
				log.Printf("Failed to initialize game '%s'. Error: %s\n", m.Game, initializationError)
				fmt.Fprintf(w, jsonError("Could not initialize game."))
				return
			}
			s.games = append(s.games, game)
			log.Printf("Created new game '%s'\n", m.Game)
		}

		player := game.GetPlayerByGoogleID(m.Player)
		// add player if it doesn't exist
		if player == nil {
			if len(game.Players) >= lib.MaxPlayers {
				log.Printf("Attempting to add player '%s' to full game '%s'\n", m.Player, m.Game)
				fmt.Fprintf(w, jsonError("This game is now full."))
				return
			}
			addError := game.AddPlayer(m.Player, sanitizeAndTrim(authResponse.GetGivenName(), lib.MaxPlayerNameLength, true))
			if addError != "" {
				log.Printf("Error adding player '%s' to game '%s'. Error: %s\n", m.Player, m.Game, addError)
				fmt.Fprintf(w, jsonError("Unable to join this game."))
				return
			}
			log.Printf("Added player '%s' to game '%s'\n", m.Player, m.Game)
		}
	}

	if command == "status" {
		if game == nil {
			return
		}
	}

	if game == nil {
		log.Printf("Attempting to make a move on a nonexistent game '%s'\n", m.Game)
		fmt.Fprintf(w, jsonError("The game you're attempting to play no longer exists."))
		return
	}

	player := game.GetPlayerByGoogleID(m.Player)

	if player == nil {
		log.Printf("Attempting to make a move with nonexistent player '%s'\n", m.Player)
		fmt.Fprintf(w, jsonError("You're not a member of this game."))
		return
	}

	if command == "start" {
		if game.State != lib.StateNotStarted {
			log.Printf("Attempting to start already started game '%s'\n", m.Game)
			fmt.Fprintf(w, jsonError("This game has already started."))
			return
		}
		var startError = game.Start()
		if startError != "" {
			log.Printf("Failed to start game '%s'. Error: %s\n", m.Game, startError)
			fmt.Fprintf(w, jsonError("Could not start game."))
			return
		}
		log.Printf("Started game '%s'\n", m.Game)
	}

	if command == "move" {
		if m.MoveType == lib.MoveHint && game.Hints <= 0 {
			fmt.Fprintf(w, jsonError("There are no hints left. Discard to earn more hints."))
			return
		}
		var processError = game.ProcessMove(m)
		if processError != "" {
			log.Printf("Failed to process move for game '%s'. Error: %s\n", m.Game, processError)
			fmt.Fprintf(w, jsonError("Could not process move."))
			return
		}
		log.Printf("Processed move by player '%s' in game '%s'\n", m.Player, m.Game)
	}

	encodedGame, err := lib.EncodeGame(game.CreateState(m.Player))
	if err != "" {
		log.Printf("Failed to encode game '%s'. Error: %s\n", m.Game, err)
		fmt.Fprintf(w, jsonError("Could not transmit game state to client."))
		return
	}
	fmt.Fprintf(w, encodedGame)
}

func jsonError(err string) string {
	return "{\"error\":\"" + strings.Replace(err, "\"", "\\\"", -1) + "\"}"
}

func sanitizeAndTrim(text string, limit int, oneword bool) string {
	re := regexp.MustCompile("[^A-Za-z0-9 _!,\\.-]+")
	text = re.ReplaceAllString(text, "")
	if oneword && strings.Index(text, " ") > -1 {
		text = text[:strings.Index(text, " ")]
	}
	if len(text) > limit {
		return text[:limit]
	}
	return text
}

type Server struct {
	games           []*lib.Game
	httpServer      bool
	clientDirectory string
}

func main() {
	rand.Seed(time.Now().UTC().UnixNano())

	// initialize server
	s := Server{}
	s.games = make([]*lib.Game, 0, lib.MaxConcurrentGames)

	// listen for connections
	httpServer := flag.Bool("http-server", false, "Whether to also serve HTTP responses outside API calls.")
	clientDirectory := flag.String("client-directory", lib.ClientDirectory, "Directory to serve HTTP responses from (fireworks-client directory)")
	port := flag.Int("port", lib.Port, "Port to listen for connections from client.")
	flag.Parse()

	s.httpServer = *httpServer
	s.clientDirectory = *clientDirectory
	http.HandleFunc("/", s.handler)
	log.Fatal(http.ListenAndServe(":"+strconv.Itoa(*port), nil))

}
<commit_msg>Improved error message for authentication failures.<commit_after>package main

import (
	"flag"
	"fmt"
	"github.com/rschoen/fireworks-server/lib"
	"log"
	"math/rand"
	"net/http"
	"regexp"
	"strconv"
	"strings"
	"time"
)

func (s *Server) handler(w http.ResponseWriter, r *http.Request) {
	// allow requests to come from anywhere, since clients can be wherever
	w.Header().Set("Access-Control-Allow-Origin", "*")

	// serve client HTTP responses, if it's turned on
	if len(r.URL.Path) < 5 || r.URL.Path[1:5] != "api/" {
		if s.httpServer == true {
			http.FileServer(http.Dir(s.clientDirectory)).ServeHTTP(w, r)
		}
		return
	}

	m, err := lib.DecodeMove(r.PostFormValue("data"))
	if err != "" {
		log.Println("Discarding malformed JSON message. Error: " + err)
		fmt.Fprintf(w, jsonError("Data sent was malformed."))
		return
	}

	var command = r.URL.Path[5:]
	var game *lib.Game
	for _, ongoingGame := range s.games {
		if ongoingGame.GameID == m.Game {
			game = ongoingGame
		}
	}

	// Authenticate user
	authResponse, authError := lib.Authenticate(m.Token)
	if authError != "" {
		log.Printf("Failed to authenticate player '%s' in game '%s'. Error: %s\n", m.Player, m.Game, authError)
		fmt.Fprintf(w, jsonError("You appear to be signed out. Please refresh and try signing in again."))
		return
	}
	if authResponse.GetGoogleID() != m.Player {
		log.Printf("Authenticated player '%s' submitted move as player '%s' in game '%s'.", authResponse.GetGoogleID(), m.Player, m.Game, m.Game)
		fmt.Fprintf(w, jsonError("Authenticated as a different user."))
		return
	}

	if command == "list" {
		list := lib.GamesList{}
		for i, _ := range s.games {
			playerList := ""
			addGame := false
			for player, _ := range s.games[i].Players {
				playerList += s.games[i].Players[player].Name + ", "
				if s.games[i].Players[player].GoogleID == m.Player {
					addGame = true
				}
			}
			if playerList != "" {
				playerList = playerList[:len(playerList)-2]
			}
			game := lib.MinimalGame{Name: s.games[i].GameID, Players: playerList}

			if addGame {
				list.PlayersGames = append(list.PlayersGames, game)
			} else if s.games[i].State == lib.StateNotStarted && len(s.games[i].Players) < lib.MaxPlayers {
				list.OpenGames = append(list.OpenGames, game)
			}
		}

		encodedList, err := lib.EncodeList(list)
		if err != "" {
			log.Printf("Failed to encode game list. Error: %s\n", err)
			fmt.Fprintf(w, jsonError("Could not transmit game list to client."))
			return
		}
		fmt.Fprintf(w, encodedList)
		return
	}

	if command == "join" {
		// create game if it doesn't exist
		if game == nil {
			game = new(lib.Game)
			game.GameID = sanitizeAndTrim(m.Game, lib.MaxGameNameLength, false)
			var initializationError = game.Initialize()
			if initializationError != "" {
				log.Printf("Failed to initialize game '%s'. Error: %s\n", m.Game, initializationError)
				fmt.Fprintf(w, jsonError("Could not initialize game."))
				return
			}
			s.games = append(s.games, game)
			log.Printf("Created new game '%s'\n", m.Game)
		}

		player := game.GetPlayerByGoogleID(m.Player)
		// add player if it doesn't exist
		if player == nil {
			if len(game.Players) >= lib.MaxPlayers {
				log.Printf("Attempting to add player '%s' to full game '%s'\n", m.Player, m.Game)
				fmt.Fprintf(w, jsonError("This game is now full."))
				return
			}
			addError := game.AddPlayer(m.Player, sanitizeAndTrim(authResponse.GetGivenName(), lib.MaxPlayerNameLength, true))
			if addError != "" {
				log.Printf("Error adding player '%s' to game '%s'. Error: %s\n", m.Player, m.Game, addError)
				fmt.Fprintf(w, jsonError("Unable to join this game."))
				return
			}
			log.Printf("Added player '%s' to game '%s'\n", m.Player, m.Game)
		}
	}

	if command == "status" {
		if game == nil {
			return
		}
	}

	if game == nil {
		log.Printf("Attempting to make a move on a nonexistent game '%s'\n", m.Game)
		fmt.Fprintf(w, jsonError("The game you're attempting to play no longer exists."))
		return
	}

	player := game.GetPlayerByGoogleID(m.Player)

	if player == nil {
		log.Printf("Attempting to make a move with nonexistent player '%s'\n", m.Player)
		fmt.Fprintf(w, jsonError("You're not a member of this game."))
		return
	}

	if command == "start" {
		if game.State != lib.StateNotStarted {
			log.Printf("Attempting to start already started game '%s'\n", m.Game)
			fmt.Fprintf(w, jsonError("This game has already started."))
			return
		}
		var startError = game.Start()
		if startError != "" {
			log.Printf("Failed to start game '%s'. Error: %s\n", m.Game, startError)
			fmt.Fprintf(w, jsonError("Could not start game."))
			return
		}
		log.Printf("Started game '%s'\n", m.Game)
	}

	if command == "move" {
		if m.MoveType == lib.MoveHint && game.Hints <= 0 {
			fmt.Fprintf(w, jsonError("There are no hints left. Discard to earn more hints."))
			return
		}
		var processError = game.ProcessMove(m)
		if processError != "" {
			log.Printf("Failed to process move for game '%s'. Error: %s\n", m.Game, processError)
			fmt.Fprintf(w, jsonError("Could not process move."))
			return
		}
		log.Printf("Processed move by player '%s' in game '%s'\n", m.Player, m.Game)
	}

	encodedGame, err := lib.EncodeGame(game.CreateState(m.Player))
	if err != "" {
		log.Printf("Failed to encode game '%s'. Error: %s\n", m.Game, err)
		fmt.Fprintf(w, jsonError("Could not transmit game state to client."))
		return
	}
	fmt.Fprintf(w, encodedGame)
}

func jsonError(err string) string {
	return "{\"error\":\"" + strings.Replace(err, "\"", "\\\"", -1) + "\"}"
}

func sanitizeAndTrim(text string, limit int, oneword bool) string {
	re := regexp.MustCompile("[^A-Za-z0-9 _!,\\.-]+")
	text = re.ReplaceAllString(text, "")
	if oneword && strings.Index(text, " ") > -1 {
		text = text[:strings.Index(text, " ")]
	}
	if len(text) > limit {
		return text[:limit]
	}
	return text
}

type Server struct {
	games           []*lib.Game
	httpServer      bool
	clientDirectory string
}

func main() {
	rand.Seed(time.Now().UTC().UnixNano())

	// initialize server
	s := Server{}
	s.games = make([]*lib.Game, 0, lib.MaxConcurrentGames)

	// listen for connections
	httpServer := flag.Bool("http-server", false, "Whether to also serve HTTP responses outside API calls.")
	clientDirectory := flag.String("client-directory", lib.ClientDirectory, "Directory to serve HTTP responses from (fireworks-client directory)")
	port := flag.Int("port", lib.Port, "Port to listen for connections from client.")
	flag.Parse()

	s.httpServer = *httpServer
	s.clientDirectory = *clientDirectory
	http.HandleFunc("/", s.handler)
	log.Fatal(http.ListenAndServe(":"+strconv.Itoa(*port), nil))

}
<|endoftext|>


<commit_before>// +build fvtests

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

package fv_test

import (
	"strconv"

	. "github.com/onsi/ginkgo"
	. "github.com/onsi/gomega"
	"github.com/projectcalico/felix/fv/containers"
	"github.com/projectcalico/felix/fv/workload"
	"github.com/projectcalico/libcalico-go/lib/api"
	"github.com/projectcalico/libcalico-go/lib/client"
	"github.com/projectcalico/libcalico-go/lib/numorstring"
)

// Setup for planned further FV tests:
//
//     | +-----------+ +-----------+ |  | +-----------+ +-----------+ |
//     | | service A | | service B | |  | | service C | | service D | |
//     | | 10.65.0.2 | | 10.65.0.3 | |  | | 10.65.0.4 | | 10.65.0.5 | |
//     | | port 9002 | | port 9003 | |  | | port 9004 | | port 9005 | |
//     | | np 109002 | | port 9003 | |  | | port 9004 | | port 9005 | |
//     | +-----------+ +-----------+ |  | +-----------+ +-----------+ |
//     +-----------------------------+  +-----------------------------+

var _ = Context("with initialized Felix, etcd datastore, 2 workloads", func() {

	var (
		etcd   *containers.Container
		felix  *containers.Container
		client *client.Client
		w      [2]*workload.Workload
	)

	BeforeEach(func() {

		etcd = RunEtcd()

		client = GetEtcdClient(etcd.IP)
		err := client.EnsureInitialized()
		Expect(err).NotTo(HaveOccurred())

		felix = RunFelix(etcd.IP)

		felixNode := api.NewNode()
		felixNode.Metadata.Name = felix.Hostname
		_, err = client.Nodes().Create(felixNode)
		Expect(err).NotTo(HaveOccurred())

		// Install a default profile that allows all ingress and egress, in the absence of any Policy.
		defaultProfile := api.NewProfile()
		defaultProfile.Metadata.Name = "default"
		defaultProfile.Metadata.Tags = []string{"default"}
		defaultProfile.Spec.EgressRules = []api.Rule{{Action: "allow"}}
		defaultProfile.Spec.IngressRules = []api.Rule{{Action: "allow"}}
		_, err = client.Profiles().Create(defaultProfile)
		Expect(err).NotTo(HaveOccurred())

		// Create workloads, using that profile.
		for ii := range w {
			iiStr := strconv.Itoa(ii)
			w[ii] = workload.Run(felix, "w"+iiStr, "cali1"+iiStr, "10.65.0.1"+iiStr, "8055")
			w[ii].Configure(client)
		}

		// We will use the etcd container to model an external client trying to connect into
		// workloads on a host.  Create a route in the etcd container for the workload CIDR.
		etcd.Exec("ip", "r", "add", "10.65.0.0/24", "via", felix.IP)
	})

	AfterEach(func() {

		if CurrentGinkgoTestDescription().Failed {
			felix.Exec("iptables-save", "-c")
			felix.Exec("ip", "r")
		}

		for ii := range w {
			w[ii].Stop()
		}
		felix.Stop()

		if CurrentGinkgoTestDescription().Failed {
			etcd.Exec("etcdctl", "ls", "--recursive", "/")
		}
		etcd.Stop()
	})

	Context("with node port DNATs", func() {

		BeforeEach(func() {
			felix.Exec(
				"iptables", "-t", "nat",
				"-A", "PREROUTING",
				"-p", "tcp",
				"-d", "10.65.0.10", "--dport", "32010",
				"-j", "DNAT", "--to", "10.65.0.10:8055",
			)
			felix.Exec(
				"iptables", "-t", "nat",
				"-A", "PREROUTING",
				"-p", "tcp",
				"-d", "10.65.0.11", "--dport", "32011",
				"-j", "DNAT", "--to", "10.65.0.11:8055",
			)
		})

		It("everyone can connect to node ports", func() {
			Expect(w[0]).To(HaveConnectivityTo(w[1], 32011))
			Expect(w[1]).To(HaveConnectivityTo(w[0], 32010))
			Expect(etcd).To(HaveConnectivityTo(w[1], 32011))
			Expect(etcd).To(HaveConnectivityTo(w[0], 32010))
		})

		Context("with pre-DNAT policy to prevent access from outside", func() {

			BeforeEach(func() {
				policy := api.NewPolicy()
				policy.Metadata.Name = "deny-ingress"
				order := float64(20)
				policy.Spec.Order = &order
				policy.Spec.PreDNAT = true
				policy.Spec.IngressRules = []api.Rule{{Action: "deny"}}
				policy.Spec.Selector = "has(host-endpoint)"
				_, err := client.Policies().Create(policy)
				Expect(err).NotTo(HaveOccurred())

				hostEp := api.NewHostEndpoint()
				hostEp.Metadata.Name = "felix-eth0"
				hostEp.Metadata.Node = felix.Hostname
				hostEp.Metadata.Labels = map[string]string{"host-endpoint": "true"}
				hostEp.Spec.InterfaceName = "eth0"
				_, err = client.HostEndpoints().Create(hostEp)
				Expect(err).NotTo(HaveOccurred())
			})

			It("etcd cannot connect", func() {
				Expect(w[0]).To(HaveConnectivityTo(w[1], 32011))
				Expect(w[1]).To(HaveConnectivityTo(w[0], 32010))
				Expect(etcd).NotTo(HaveConnectivityTo(w[1], 32011))
				Expect(etcd).NotTo(HaveConnectivityTo(w[0], 32010))
			})

			Context("with pre-DNAT policy to open pinhole to 32010", func() {

				BeforeEach(func() {
					policy := api.NewPolicy()
					policy.Metadata.Name = "allow-ingress-32010"
					order := float64(10)
					policy.Spec.Order = &order
					policy.Spec.PreDNAT = true
					protocol := numorstring.ProtocolFromString("tcp")
					ports, _ := numorstring.PortFromRange(1, 65535)   // pass
					ports, _ = numorstring.PortFromRange(1, 10000)    // pass
					ports, _ = numorstring.PortFromRange(8000, 8055)  // pass
					ports, _ = numorstring.PortFromRange(8056, 10000) // fail
					ports = numorstring.SinglePort(8055)
					policy.Spec.IngressRules = []api.Rule{{
						Action:   "allow",
						Protocol: &protocol,
						Destination: api.EntityRule{Ports: []numorstring.Port{
							ports,
						}},
					}}
					policy.Spec.Selector = "has(host-endpoint)"
					_, err := client.Policies().Create(policy)
					Expect(err).NotTo(HaveOccurred())
				})

				// Pending because currently fails - investigation needed.
				PIt("etcd can connect to 32010 but not 32011", func() {
					Expect(w[0]).To(HaveConnectivityTo(w[1], 32011))
					Expect(w[1]).To(HaveConnectivityTo(w[0], 32010))
					//Expect(etcd).NotTo(HaveConnectivityTo(w[1], 32011))
					var success bool
					var err error
					success, err = HaveConnectivityTo(w[0], 32010).Match(etcd)
					success, err = HaveConnectivityTo(w[0], 32010).Match(etcd)
					success, err = HaveConnectivityTo(w[0], 32010).Match(etcd)
					success, err = HaveConnectivityTo(w[0], 32010).Match(etcd)
					success, err = HaveConnectivityTo(w[0], 32010).Match(etcd)
					success, err = HaveConnectivityTo(w[0], 32010).Match(etcd)
					success, err = HaveConnectivityTo(w[0], 32010).Match(etcd)
					success, err = HaveConnectivityTo(w[0], 32010).Match(etcd)
					success, err = HaveConnectivityTo(w[0], 32010).Match(etcd)
					success, err = HaveConnectivityTo(w[0], 32010).Match(etcd)
					Expect(success).To(BeTrue())
					Expect(err).NotTo(HaveOccurred())
					Expect(etcd).NotTo(HaveConnectivityTo(w[0], 32010))
				})
			})
		})
	})
})
<commit_msg>Review markups<commit_after>// +build fvtests

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

package fv_test

import (
	"strconv"

	. "github.com/onsi/ginkgo"
	. "github.com/onsi/gomega"
	"github.com/projectcalico/felix/fv/containers"
	"github.com/projectcalico/felix/fv/workload"
	"github.com/projectcalico/libcalico-go/lib/api"
	"github.com/projectcalico/libcalico-go/lib/client"
	"github.com/projectcalico/libcalico-go/lib/numorstring"
)

// Setup for planned further FV tests:
//
//     | +-----------+ +-----------+ |  | +-----------+ +-----------+ |
//     | | service A | | service B | |  | | service C | | service D | |
//     | | 10.65.0.2 | | 10.65.0.3 | |  | | 10.65.0.4 | | 10.65.0.5 | |
//     | | port 9002 | | port 9003 | |  | | port 9004 | | port 9005 | |
//     | | np 109002 | | port 9003 | |  | | port 9004 | | port 9005 | |
//     | +-----------+ +-----------+ |  | +-----------+ +-----------+ |
//     +-----------------------------+  +-----------------------------+

var _ = Context("with initialized Felix, etcd datastore, 2 workloads", func() {

	var (
		etcd   *containers.Container
		felix  *containers.Container
		client *client.Client
		w      [2]*workload.Workload
	)

	BeforeEach(func() {

		etcd = RunEtcd()

		client = GetEtcdClient(etcd.IP)
		err := client.EnsureInitialized()
		Expect(err).NotTo(HaveOccurred())

		felix = RunFelix(etcd.IP)

		felixNode := api.NewNode()
		felixNode.Metadata.Name = felix.Hostname
		_, err = client.Nodes().Create(felixNode)
		Expect(err).NotTo(HaveOccurred())

		// Install a default profile that allows all ingress and egress, in the absence of any Policy.
		defaultProfile := api.NewProfile()
		defaultProfile.Metadata.Name = "default"
		defaultProfile.Metadata.Tags = []string{"default"}
		defaultProfile.Spec.EgressRules = []api.Rule{{Action: "allow"}}
		defaultProfile.Spec.IngressRules = []api.Rule{{Action: "allow"}}
		_, err = client.Profiles().Create(defaultProfile)
		Expect(err).NotTo(HaveOccurred())

		// Create workloads, using that profile.
		for ii := range w {
			iiStr := strconv.Itoa(ii)
			w[ii] = workload.Run(felix, "w"+iiStr, "cali1"+iiStr, "10.65.0.1"+iiStr, "8055")
			w[ii].Configure(client)
		}

		// We will use the etcd container to model an external client trying to connect into
		// workloads on a host.  Create a route in the etcd container for the workload CIDR.
		etcd.Exec("ip", "r", "add", "10.65.0.0/24", "via", felix.IP)
	})

	AfterEach(func() {

		if CurrentGinkgoTestDescription().Failed {
			felix.Exec("iptables-save", "-c")
			felix.Exec("ip", "r")
		}

		for ii := range w {
			w[ii].Stop()
		}
		felix.Stop()

		if CurrentGinkgoTestDescription().Failed {
			etcd.Exec("etcdctl", "ls", "--recursive", "/")
		}
		etcd.Stop()
	})

	Context("with node port DNATs", func() {

		BeforeEach(func() {
			felix.Exec(
				"iptables", "-t", "nat",
				"-A", "PREROUTING",
				"-p", "tcp",
				"-d", "10.65.0.10", "--dport", "32010",
				"-j", "DNAT", "--to", "10.65.0.10:8055",
			)
			felix.Exec(
				"iptables", "-t", "nat",
				"-A", "PREROUTING",
				"-p", "tcp",
				"-d", "10.65.0.11", "--dport", "32011",
				"-j", "DNAT", "--to", "10.65.0.11:8055",
			)
		})

		It("everyone can connect to node ports", func() {
			cc := &workload.ConnectivityChecker{}
			cc.ExpectSome(w[0], w[1], 32011)
			cc.ExpectSome(w[1], w[0], 32010)
			cc.ExpectSome(etcd, w[1], 32011)
			cc.ExpectSome(etcd, w[0], 32010)
			Eventually(cc.ActualConnectivity, "10s", "100ms").Should(Equal(cc.ExpectedConnectivity()))
		})

		Context("with pre-DNAT policy to prevent access from outside", func() {

			BeforeEach(func() {
				policy := api.NewPolicy()
				policy.Metadata.Name = "deny-ingress"
				order := float64(20)
				policy.Spec.Order = &order
				policy.Spec.PreDNAT = true
				policy.Spec.IngressRules = []api.Rule{{Action: "deny"}}
				policy.Spec.Selector = "has(host-endpoint)"
				_, err := client.Policies().Create(policy)
				Expect(err).NotTo(HaveOccurred())

				hostEp := api.NewHostEndpoint()
				hostEp.Metadata.Name = "felix-eth0"
				hostEp.Metadata.Node = felix.Hostname
				hostEp.Metadata.Labels = map[string]string{"host-endpoint": "true"}
				hostEp.Spec.InterfaceName = "eth0"
				_, err = client.HostEndpoints().Create(hostEp)
				Expect(err).NotTo(HaveOccurred())
			})

			It("etcd cannot connect", func() {
				cc := &workload.ConnectivityChecker{}
				cc.ExpectSome(w[0], w[1], 32011)
				cc.ExpectSome(w[1], w[0], 32010)
				cc.ExpectNone(etcd, w[1], 32011)
				cc.ExpectNone(etcd, w[0], 32010)
				Eventually(cc.ActualConnectivity, "10s", "100ms").Should(Equal(cc.ExpectedConnectivity()))
			})

			Context("with pre-DNAT policy to open pinhole to 32010", func() {

				BeforeEach(func() {
					policy := api.NewPolicy()
					policy.Metadata.Name = "allow-ingress-32010"
					order := float64(10)
					policy.Spec.Order = &order
					policy.Spec.PreDNAT = true
					protocol := numorstring.ProtocolFromString("tcp")
					ports := numorstring.SinglePort(32010)
					policy.Spec.IngressRules = []api.Rule{{
						Action:   "allow",
						Protocol: &protocol,
						Destination: api.EntityRule{Ports: []numorstring.Port{
							ports,
						}},
					}}
					policy.Spec.Selector = "has(host-endpoint)"
					_, err := client.Policies().Create(policy)
					Expect(err).NotTo(HaveOccurred())
				})

				It("etcd can connect to 32010 but not 32011", func() {
					cc := &workload.ConnectivityChecker{}
					cc.ExpectSome(w[0], w[1], 32011)
					cc.ExpectSome(w[1], w[0], 32010)
					cc.ExpectNone(etcd, w[1], 32011)
					cc.ExpectSome(etcd, w[0], 32010)
					Eventually(cc.ActualConnectivity, "10s", "100ms").Should(Equal(cc.ExpectedConnectivity()))
				})
			})

			Context("with pre-DNAT policy to open pinhole to 8055", func() {

				BeforeEach(func() {
					policy := api.NewPolicy()
					policy.Metadata.Name = "allow-ingress-8055"
					order := float64(10)
					policy.Spec.Order = &order
					policy.Spec.PreDNAT = true
					protocol := numorstring.ProtocolFromString("tcp")
					ports := numorstring.SinglePort(8055)
					policy.Spec.IngressRules = []api.Rule{{
						Action:   "allow",
						Protocol: &protocol,
						Destination: api.EntityRule{Ports: []numorstring.Port{
							ports,
						}},
					}}
					policy.Spec.Selector = "has(host-endpoint)"
					_, err := client.Policies().Create(policy)
					Expect(err).NotTo(HaveOccurred())
				})

				It("etcd cannot connect", func() {
					cc := &workload.ConnectivityChecker{}
					cc.ExpectSome(w[0], w[1], 32011)
					cc.ExpectSome(w[1], w[0], 32010)
					cc.ExpectNone(etcd, w[1], 32011)
					cc.ExpectNone(etcd, w[0], 32010)
					Eventually(cc.ActualConnectivity, "10s", "100ms").Should(Equal(cc.ExpectedConnectivity()))
				})
			})
		})
	})
})
<|endoftext|>


<commit_before>/*
Copyright 2017 The Kubernetes Authors.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/

package gce

import (
	"flag"
	"fmt"
	"net"
	"strings"
	"time"

	"github.com/golang/glog"

	"k8s.io/api/core/v1"
	"k8s.io/kubernetes/pkg/cloudprovider"
	netsets "k8s.io/kubernetes/pkg/util/net/sets"
)

type cidrs struct {
	ipn   netsets.IPNet
	isSet bool
}

var (
	lbSrcRngsFlag cidrs
)

func newLoadBalancerMetricContext(request, region string) *metricContext {
	return &metricContext{
		start:      time.Now(),
		attributes: []string{"loadbalancer_" + request, region, unusedMetricLabel},
	}
}

type lbScheme string

const (
	schemeExternal lbScheme = "EXTERNAL"
	schemeInternal lbScheme = "INTERNAL"
)

func init() {
	var err error
	// LB L7 proxies and all L3/4/7 health checkers have client addresses within these known CIDRs.
	lbSrcRngsFlag.ipn, err = netsets.ParseIPNets([]string{"130.211.0.0/22", "35.191.0.0/16", "209.85.152.0/22", "209.85.204.0/22"}...)
	if err != nil {
		panic("Incorrect default GCE L7 source ranges")
	}

	flag.Var(&lbSrcRngsFlag, "cloud-provider-gce-lb-src-cidrs", "CIDRS opened in GCE firewall for LB traffic proxy & health checks")
}

// String is the method to format the flag's value, part of the flag.Value interface.
func (c *cidrs) String() string {
	return strings.Join(c.ipn.StringSlice(), ",")
}

// Set supports a value of CSV or the flag repeated multiple times
func (c *cidrs) Set(value string) error {
	// On first Set(), clear the original defaults
	if !c.isSet {
		c.isSet = true
		c.ipn = make(netsets.IPNet)
	} else {
		return fmt.Errorf("GCE LB CIDRS have already been set")
	}

	for _, cidr := range strings.Split(value, ",") {
		_, ipnet, err := net.ParseCIDR(cidr)
		if err != nil {
			return err
		}

		c.ipn.Insert(ipnet)
	}
	return nil
}

// LoadBalancerSrcRanges contains the ranges of ips used by the GCE load balancers (l4 & L7)
// for proxying client requests and performing health checks.
func LoadBalancerSrcRanges() []string {
	return lbSrcRngsFlag.ipn.StringSlice()
}

// GetLoadBalancer is an implementation of LoadBalancer.GetLoadBalancer
func (gce *GCECloud) GetLoadBalancer(clusterName string, svc *v1.Service) (*v1.LoadBalancerStatus, bool, error) {
	loadBalancerName := cloudprovider.GetLoadBalancerName(svc)
	fwd, err := gce.GetRegionForwardingRule(loadBalancerName, gce.region)
	if err == nil {
		status := &v1.LoadBalancerStatus{}
		status.Ingress = []v1.LoadBalancerIngress{{IP: fwd.IPAddress}}

		return status, true, nil
	}
	return nil, false, ignoreNotFound(err)
}

// EnsureLoadBalancer is an implementation of LoadBalancer.EnsureLoadBalancer.
func (gce *GCECloud) EnsureLoadBalancer(clusterName string, svc *v1.Service, nodes []*v1.Node) (*v1.LoadBalancerStatus, error) {
	loadBalancerName := cloudprovider.GetLoadBalancerName(svc)
	desiredScheme := getSvcScheme(svc)
	clusterID, err := gce.ClusterID.GetID()
	if err != nil {
		return nil, err
	}

	glog.V(4).Infof("EnsureLoadBalancer(%v, %v, %v, %v, %v): ensure %v loadbalancer", clusterName, svc.Namespace, svc.Name, loadBalancerName, gce.region, desiredScheme)

	existingFwdRule, err := gce.GetRegionForwardingRule(loadBalancerName, gce.region)
	if err != nil && !isNotFound(err) {
		return nil, err
	}

	if existingFwdRule != nil {
		existingScheme := lbScheme(strings.ToUpper(existingFwdRule.LoadBalancingScheme))

		// If the loadbalancer type changes between INTERNAL and EXTERNAL, the old load balancer should be deleted.
		if existingScheme != desiredScheme {
			glog.V(4).Infof("EnsureLoadBalancer(%v, %v, %v, %v, %v): deleting existing %v loadbalancer", clusterName, svc.Namespace, svc.Name, loadBalancerName, gce.region, existingScheme)
			switch existingScheme {
			case schemeInternal:
				err = gce.ensureInternalLoadBalancerDeleted(clusterName, clusterID, svc)
			default:
				err = gce.ensureExternalLoadBalancerDeleted(clusterName, clusterID, svc)
			}
			glog.V(4).Infof("EnsureLoadBalancer(%v, %v, %v, %v, %v): done deleting existing %v loadbalancer. err: %v", clusterName, svc.Namespace, svc.Name, loadBalancerName, gce.region, existingScheme, err)
			if err != nil {
				return nil, err
			}
		}
	}

	var status *v1.LoadBalancerStatus
	switch desiredScheme {
	case schemeInternal:
		status, err = gce.ensureInternalLoadBalancer(clusterName, clusterID, svc, existingFwdRule, nodes)
	default:
		status, err = gce.ensureExternalLoadBalancer(clusterName, clusterID, svc, existingFwdRule, nodes)
	}
	glog.V(4).Infof("EnsureLoadBalancer(%v, %v, %v, %v, %v): done ensuring loadbalancer. err: %v", clusterName, svc.Namespace, svc.Name, loadBalancerName, gce.region, err)
	return status, err
}

// UpdateLoadBalancer is an implementation of LoadBalancer.UpdateLoadBalancer.
func (gce *GCECloud) UpdateLoadBalancer(clusterName string, svc *v1.Service, nodes []*v1.Node) error {
	loadBalancerName := cloudprovider.GetLoadBalancerName(svc)
	scheme := getSvcScheme(svc)
	clusterID, err := gce.ClusterID.GetID()
	if err != nil {
		return err
	}

	glog.V(4).Infof("UpdateLoadBalancer(%v, %v, %v, %v, %v): updating with %d nodes", clusterName, svc.Namespace, svc.Name, loadBalancerName, gce.region, len(nodes))

	switch scheme {
	case schemeInternal:
		err = gce.updateInternalLoadBalancer(clusterName, clusterID, svc, nodes)
	default:
		err = gce.updateExternalLoadBalancer(clusterName, svc, nodes)
	}
	glog.V(4).Infof("UpdateLoadBalancer(%v, %v, %v, %v, %v): done updating. err: %v", clusterName, svc.Namespace, svc.Name, loadBalancerName, gce.region, err)
	return err
}

// EnsureLoadBalancerDeleted is an implementation of LoadBalancer.EnsureLoadBalancerDeleted.
func (gce *GCECloud) EnsureLoadBalancerDeleted(clusterName string, svc *v1.Service) error {
	loadBalancerName := cloudprovider.GetLoadBalancerName(svc)
	scheme := getSvcScheme(svc)
	clusterID, err := gce.ClusterID.GetID()
	if err != nil {
		return err
	}

	glog.V(4).Infof("EnsureLoadBalancerDeleted(%v, %v, %v, %v, %v): deleting loadbalancer", clusterName, svc.Namespace, svc.Name, loadBalancerName, gce.region)

	switch scheme {
	case schemeInternal:
		err = gce.ensureInternalLoadBalancerDeleted(clusterName, clusterID, svc)
	default:
		err = gce.ensureExternalLoadBalancerDeleted(clusterName, clusterID, svc)
	}
	glog.V(4).Infof("EnsureLoadBalancerDeleted(%v, %v, %v, %v, %v): done deleting loadbalancer. err: %v", clusterName, svc.Namespace, svc.Name, loadBalancerName, gce.region, err)
	return err
}

func getSvcScheme(svc *v1.Service) lbScheme {
	if typ, ok := GetLoadBalancerAnnotationType(svc); ok && typ == LBTypeInternal {
		return schemeInternal
	}
	return schemeExternal
}
<commit_msg>fix minor typo<commit_after>/*
Copyright 2017 The Kubernetes Authors.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/

package gce

import (
	"flag"
	"fmt"
	"net"
	"strings"
	"time"

	"github.com/golang/glog"

	"k8s.io/api/core/v1"
	"k8s.io/kubernetes/pkg/cloudprovider"
	netsets "k8s.io/kubernetes/pkg/util/net/sets"
)

type cidrs struct {
	ipn   netsets.IPNet
	isSet bool
}

var (
	lbSrcRngsFlag cidrs
)

func newLoadBalancerMetricContext(request, region string) *metricContext {
	return &metricContext{
		start:      time.Now(),
		attributes: []string{"loadbalancer_" + request, region, unusedMetricLabel},
	}
}

type lbScheme string

const (
	schemeExternal lbScheme = "EXTERNAL"
	schemeInternal lbScheme = "INTERNAL"
)

func init() {
	var err error
	// LB L7 proxies and all L3/4/7 health checkers have client addresses within these known CIDRs.
	lbSrcRngsFlag.ipn, err = netsets.ParseIPNets([]string{"130.211.0.0/22", "35.191.0.0/16", "209.85.152.0/22", "209.85.204.0/22"}...)
	if err != nil {
		panic("Incorrect default GCE L7 source ranges")
	}

	flag.Var(&lbSrcRngsFlag, "cloud-provider-gce-lb-src-cidrs", "CIDRs opened in GCE firewall for LB traffic proxy & health checks")
}

// String is the method to format the flag's value, part of the flag.Value interface.
func (c *cidrs) String() string {
	return strings.Join(c.ipn.StringSlice(), ",")
}

// Set supports a value of CSV or the flag repeated multiple times
func (c *cidrs) Set(value string) error {
	// On first Set(), clear the original defaults
	if !c.isSet {
		c.isSet = true
		c.ipn = make(netsets.IPNet)
	} else {
		return fmt.Errorf("GCE LB CIDRs have already been set")
	}

	for _, cidr := range strings.Split(value, ",") {
		_, ipnet, err := net.ParseCIDR(cidr)
		if err != nil {
			return err
		}

		c.ipn.Insert(ipnet)
	}
	return nil
}

// LoadBalancerSrcRanges contains the ranges of ips used by the GCE load balancers (l4 & L7)
// for proxying client requests and performing health checks.
func LoadBalancerSrcRanges() []string {
	return lbSrcRngsFlag.ipn.StringSlice()
}

// GetLoadBalancer is an implementation of LoadBalancer.GetLoadBalancer
func (gce *GCECloud) GetLoadBalancer(clusterName string, svc *v1.Service) (*v1.LoadBalancerStatus, bool, error) {
	loadBalancerName := cloudprovider.GetLoadBalancerName(svc)
	fwd, err := gce.GetRegionForwardingRule(loadBalancerName, gce.region)
	if err == nil {
		status := &v1.LoadBalancerStatus{}
		status.Ingress = []v1.LoadBalancerIngress{{IP: fwd.IPAddress}}

		return status, true, nil
	}
	return nil, false, ignoreNotFound(err)
}

// EnsureLoadBalancer is an implementation of LoadBalancer.EnsureLoadBalancer.
func (gce *GCECloud) EnsureLoadBalancer(clusterName string, svc *v1.Service, nodes []*v1.Node) (*v1.LoadBalancerStatus, error) {
	loadBalancerName := cloudprovider.GetLoadBalancerName(svc)
	desiredScheme := getSvcScheme(svc)
	clusterID, err := gce.ClusterID.GetID()
	if err != nil {
		return nil, err
	}

	glog.V(4).Infof("EnsureLoadBalancer(%v, %v, %v, %v, %v): ensure %v loadbalancer", clusterName, svc.Namespace, svc.Name, loadBalancerName, gce.region, desiredScheme)

	existingFwdRule, err := gce.GetRegionForwardingRule(loadBalancerName, gce.region)
	if err != nil && !isNotFound(err) {
		return nil, err
	}

	if existingFwdRule != nil {
		existingScheme := lbScheme(strings.ToUpper(existingFwdRule.LoadBalancingScheme))

		// If the loadbalancer type changes between INTERNAL and EXTERNAL, the old load balancer should be deleted.
		if existingScheme != desiredScheme {
			glog.V(4).Infof("EnsureLoadBalancer(%v, %v, %v, %v, %v): deleting existing %v loadbalancer", clusterName, svc.Namespace, svc.Name, loadBalancerName, gce.region, existingScheme)
			switch existingScheme {
			case schemeInternal:
				err = gce.ensureInternalLoadBalancerDeleted(clusterName, clusterID, svc)
			default:
				err = gce.ensureExternalLoadBalancerDeleted(clusterName, clusterID, svc)
			}
			glog.V(4).Infof("EnsureLoadBalancer(%v, %v, %v, %v, %v): done deleting existing %v loadbalancer. err: %v", clusterName, svc.Namespace, svc.Name, loadBalancerName, gce.region, existingScheme, err)
			if err != nil {
				return nil, err
			}
		}
	}

	var status *v1.LoadBalancerStatus
	switch desiredScheme {
	case schemeInternal:
		status, err = gce.ensureInternalLoadBalancer(clusterName, clusterID, svc, existingFwdRule, nodes)
	default:
		status, err = gce.ensureExternalLoadBalancer(clusterName, clusterID, svc, existingFwdRule, nodes)
	}
	glog.V(4).Infof("EnsureLoadBalancer(%v, %v, %v, %v, %v): done ensuring loadbalancer. err: %v", clusterName, svc.Namespace, svc.Name, loadBalancerName, gce.region, err)
	return status, err
}

// UpdateLoadBalancer is an implementation of LoadBalancer.UpdateLoadBalancer.
func (gce *GCECloud) UpdateLoadBalancer(clusterName string, svc *v1.Service, nodes []*v1.Node) error {
	loadBalancerName := cloudprovider.GetLoadBalancerName(svc)
	scheme := getSvcScheme(svc)
	clusterID, err := gce.ClusterID.GetID()
	if err != nil {
		return err
	}

	glog.V(4).Infof("UpdateLoadBalancer(%v, %v, %v, %v, %v): updating with %d nodes", clusterName, svc.Namespace, svc.Name, loadBalancerName, gce.region, len(nodes))

	switch scheme {
	case schemeInternal:
		err = gce.updateInternalLoadBalancer(clusterName, clusterID, svc, nodes)
	default:
		err = gce.updateExternalLoadBalancer(clusterName, svc, nodes)
	}
	glog.V(4).Infof("UpdateLoadBalancer(%v, %v, %v, %v, %v): done updating. err: %v", clusterName, svc.Namespace, svc.Name, loadBalancerName, gce.region, err)
	return err
}

// EnsureLoadBalancerDeleted is an implementation of LoadBalancer.EnsureLoadBalancerDeleted.
func (gce *GCECloud) EnsureLoadBalancerDeleted(clusterName string, svc *v1.Service) error {
	loadBalancerName := cloudprovider.GetLoadBalancerName(svc)
	scheme := getSvcScheme(svc)
	clusterID, err := gce.ClusterID.GetID()
	if err != nil {
		return err
	}

	glog.V(4).Infof("EnsureLoadBalancerDeleted(%v, %v, %v, %v, %v): deleting loadbalancer", clusterName, svc.Namespace, svc.Name, loadBalancerName, gce.region)

	switch scheme {
	case schemeInternal:
		err = gce.ensureInternalLoadBalancerDeleted(clusterName, clusterID, svc)
	default:
		err = gce.ensureExternalLoadBalancerDeleted(clusterName, clusterID, svc)
	}
	glog.V(4).Infof("EnsureLoadBalancerDeleted(%v, %v, %v, %v, %v): done deleting loadbalancer. err: %v", clusterName, svc.Namespace, svc.Name, loadBalancerName, gce.region, err)
	return err
}

func getSvcScheme(svc *v1.Service) lbScheme {
	if typ, ok := GetLoadBalancerAnnotationType(svc); ok && typ == LBTypeInternal {
		return schemeInternal
	}
	return schemeExternal
}
<|endoftext|>


<commit_before>package generator

import (
	"fmt"
	"go/types"
	"log"
	"reflect"
	"strings"

	"golang.org/x/tools/go/packages"
	"golang.org/x/tools/imports"
)

func (f *Fake) loadPackages(c Cacher, workingDir string) error {
	log.Println("loading packages...")
	p, ok := c.Load(f.TargetPackage)
	if ok {
		f.Packages = p
		log.Printf("loaded %v packages from cache\n", len(f.Packages))
		return nil
	}
	p, err := packages.Load(&packages.Config{
		Mode:  packages.LoadSyntax,
		Dir:   workingDir,
		Tests: true,
	}, f.TargetPackage)
	if err != nil {
		return err
	}
	for i := range p {
		if len(p[i].Errors) > 0 {
			if i == 0 {
				err = p[i].Errors[0]
			}
			for j := range p[i].Errors {
				log.Printf("error loading packages: %v", strings.TrimPrefix(fmt.Sprintf("%v", p[i].Errors[j]), "-: "))
			}
		}
	}
	if err != nil {
		return err
	}
	f.Packages = p
	c.Store(f.TargetPackage, p)
	log.Printf("loaded %v packages\n", len(f.Packages))
	return nil
}

func (f *Fake) findPackage() error {
	var target *types.TypeName
	var pkg *packages.Package
	for i := range f.Packages {
		if f.Packages[i].Types == nil || f.Packages[i].Types.Scope() == nil {
			continue
		}
		pkg = f.Packages[i]
		if f.Mode == Package {
			break
		}

		raw := pkg.Types.Scope().Lookup(f.TargetName)
		if raw != nil {
			if typeName, ok := raw.(*types.TypeName); ok {
				target = typeName
				break
			}
		}
		pkg = nil
	}
	if pkg == nil {
		switch f.Mode {
		case Package:
			return fmt.Errorf("cannot find package with name: %s", f.TargetPackage)
		case InterfaceOrFunction:
			return fmt.Errorf("cannot find package with target: %s", f.TargetName)
		}
	}
	f.Target = target
	f.Package = pkg
	f.TargetPackage = imports.VendorlessPath(pkg.PkgPath)
	t := f.Imports.Add(pkg.Name, f.TargetPackage)
	f.TargetAlias = t.Alias
	if f.Mode != Package {
		f.TargetName = target.Name()
	}

	if f.Mode == InterfaceOrFunction {
		if !f.IsInterface() && !f.IsFunction() {
			return fmt.Errorf("cannot generate an fake for %s because it is not an interface or function", f.TargetName)
		}
	}

	if f.IsInterface() {
		log.Printf("Found interface with name: [%s]\n", f.TargetName)
	}
	if f.IsFunction() {
		log.Printf("Found function with name: [%s]\n", f.TargetName)
	}
	if f.Mode == Package {
		log.Printf("Found package with name: [%s]\n", f.TargetPackage)
	}
	return nil
}

// addImportsFor inspects the given type and adds imports to the fake if importable
// types are found.
func (f *Fake) addImportsFor(typ types.Type) {
	if typ == nil {
		return
	}

	switch t := typ.(type) {
	case *types.Basic:
		return
	case *types.Pointer:
		f.addImportsFor(t.Elem())
	case *types.Map:
		f.addImportsFor(t.Key())
		f.addImportsFor(t.Elem())
	case *types.Chan:
		f.addImportsFor(t.Elem())
	case *types.Named:
		if t.Obj() != nil && t.Obj().Pkg() != nil {
			f.Imports.Add(t.Obj().Pkg().Name(), t.Obj().Pkg().Path())
		}
	case *types.Slice:
		f.addImportsFor(t.Elem())
	case *types.Array:
		f.addImportsFor(t.Elem())
	case *types.Interface:
		return
	case *types.Signature:
		f.addTypesForMethod(t)
	default:
		log.Printf("!!! WARNING: Missing case for type %s\n", reflect.TypeOf(typ).String())
	}
}
<commit_msg>be specific when defining packages.Config.Mode<commit_after>package generator

import (
	"fmt"
	"go/types"
	"log"
	"reflect"
	"strings"

	"golang.org/x/tools/go/packages"
	"golang.org/x/tools/imports"
)

func (f *Fake) loadPackages(c Cacher, workingDir string) error {
	log.Println("loading packages...")
	p, ok := c.Load(f.TargetPackage)
	if ok {
		f.Packages = p
		log.Printf("loaded %v packages from cache\n", len(f.Packages))
		return nil
	}
	p, err := packages.Load(&packages.Config{
		Mode:  packages.NeedName | packages.NeedFiles | packages.NeedImports | packages.NeedDeps | packages.NeedTypes,
		Dir:   workingDir,
		Tests: true,
	}, f.TargetPackage)
	if err != nil {
		return err
	}
	for i := range p {
		if len(p[i].Errors) > 0 {
			if i == 0 {
				err = p[i].Errors[0]
			}
			for j := range p[i].Errors {
				log.Printf("error loading packages: %v", strings.TrimPrefix(fmt.Sprintf("%v", p[i].Errors[j]), "-: "))
			}
		}
	}
	if err != nil {
		return err
	}
	f.Packages = p
	c.Store(f.TargetPackage, p)
	log.Printf("loaded %v packages\n", len(f.Packages))
	return nil
}

func (f *Fake) findPackage() error {
	var target *types.TypeName
	var pkg *packages.Package
	for i := range f.Packages {
		if f.Packages[i].Types == nil || f.Packages[i].Types.Scope() == nil {
			continue
		}
		pkg = f.Packages[i]
		if f.Mode == Package {
			break
		}

		raw := pkg.Types.Scope().Lookup(f.TargetName)
		if raw != nil {
			if typeName, ok := raw.(*types.TypeName); ok {
				target = typeName
				break
			}
		}
		pkg = nil
	}
	if pkg == nil {
		switch f.Mode {
		case Package:
			return fmt.Errorf("cannot find package with name: %s", f.TargetPackage)
		case InterfaceOrFunction:
			return fmt.Errorf("cannot find package with target: %s", f.TargetName)
		}
	}
	f.Target = target
	f.Package = pkg
	f.TargetPackage = imports.VendorlessPath(pkg.PkgPath)
	t := f.Imports.Add(pkg.Name, f.TargetPackage)
	f.TargetAlias = t.Alias
	if f.Mode != Package {
		f.TargetName = target.Name()
	}

	if f.Mode == InterfaceOrFunction {
		if !f.IsInterface() && !f.IsFunction() {
			return fmt.Errorf("cannot generate an fake for %s because it is not an interface or function", f.TargetName)
		}
	}

	if f.IsInterface() {
		log.Printf("Found interface with name: [%s]\n", f.TargetName)
	}
	if f.IsFunction() {
		log.Printf("Found function with name: [%s]\n", f.TargetName)
	}
	if f.Mode == Package {
		log.Printf("Found package with name: [%s]\n", f.TargetPackage)
	}
	return nil
}

// addImportsFor inspects the given type and adds imports to the fake if importable
// types are found.
func (f *Fake) addImportsFor(typ types.Type) {
	if typ == nil {
		return
	}

	switch t := typ.(type) {
	case *types.Basic:
		return
	case *types.Pointer:
		f.addImportsFor(t.Elem())
	case *types.Map:
		f.addImportsFor(t.Key())
		f.addImportsFor(t.Elem())
	case *types.Chan:
		f.addImportsFor(t.Elem())
	case *types.Named:
		if t.Obj() != nil && t.Obj().Pkg() != nil {
			f.Imports.Add(t.Obj().Pkg().Name(), t.Obj().Pkg().Path())
		}
	case *types.Slice:
		f.addImportsFor(t.Elem())
	case *types.Array:
		f.addImportsFor(t.Elem())
	case *types.Interface:
		return
	case *types.Signature:
		f.addTypesForMethod(t)
	default:
		log.Printf("!!! WARNING: Missing case for type %s\n", reflect.TypeOf(typ).String())
	}
}
<|endoftext|>


<commit_before>
// Various functions for manipulating time objects
package util

import (
  "errors"
  "fmt"
  "github.com/jinzhu/now"
  "math"
  "strings"
  "time"
)

// Replaces the normal duration.String() function with one which formats the
// data in a much more human readable way.
func DurationToString(d time.Duration) string {
  s := ""
  hours := math.Floor(d.Hours())
  if hours > 0 {
    s += fmt.Sprintf("%v hour", hours)
  }
  if hours > 1 {
    s += "s"
  }
  minutes := int(d.Minutes()) % 60
  if minutes > 0 {
    if hours > 0 {
      s += " "
    }
    s += fmt.Sprintf("%v minute", minutes)
    if minutes > 1 {
      s += "s"
    }
  }
  seconds := int(d.Seconds()) % 60
  if seconds > 0 {
    if hours > 0 || minutes > 0 {
      s += " "
    }
    s += fmt.Sprintf("%v second", seconds)
    if seconds > 1 {
      s += "s"
    }
  }
  return s
}

func StringToDuration(durStr string) (time.Duration, error) {
  // Run the duration string through a bunch of processing to get it into a time.Duration format that can be parsed by go
  durStr = strings.Replace(durStr, " seconds", "s", -1)
  durStr = strings.Replace(durStr, " second", "s", -1)
  durStr = strings.Replace(durStr, " minutes", "m", -1)
  durStr = strings.Replace(durStr, " minute", "m", -1)
  durStr = strings.Replace(durStr, " hours", "h", -1)
  durStr = strings.Replace(durStr, " hour", "h", -1)
  durStr = strings.Replace(durStr, " ", "", -1)
  d, err := time.ParseDuration(durStr)
  if err != nil {
    return d, errors.New("Apologies, but I can't seem to parse your duration string.")
  }
  if d.Hours() < 0 || d.Minutes() < 0 || d.Seconds() < 0 {
    return d, errors.New("Apologies, but my functionality does not include the recognition of negative time.")
  }
  return d, nil
}

func StringToTime(ts string) (time.Time, error) {
  stockErr := errors.New("Apologies, but I can't seem to read the time you gave me.")
  t, err := now.Parse(ts)
  fmt.Printf("%v\n", t.String())
  if err != nil {
    return t, stockErr
  }
  return t, nil
}
<commit_msg>Added capability to parse times like '8' or '8:30' correctly based on relative offset to current time<commit_after>
// Various functions for manipulating time objects
package util

import (
  "errors"
  "fmt"
  "github.com/jinzhu/now"
  "jarvis/log"
  "math"
  "strings"
  "time"
)

// Replaces the normal duration.String() function with one which formats the
// data in a much more human readable way.
func DurationToString(d time.Duration) string {
  s := ""
  hours := math.Floor(d.Hours())
  if hours > 0 {
    s += fmt.Sprintf("%v hour", hours)
  }
  if hours > 1 {
    s += "s"
  }
  minutes := int(d.Minutes()) % 60
  if minutes > 0 {
    if hours > 0 {
      s += " "
    }
    s += fmt.Sprintf("%v minute", minutes)
    if minutes > 1 {
      s += "s"
    }
  }
  seconds := int(d.Seconds()) % 60
  if seconds > 0 {
    if hours > 0 || minutes > 0 {
      s += " "
    }
    s += fmt.Sprintf("%v second", seconds)
    if seconds > 1 {
      s += "s"
    }
  }
  return s
}

func StringToDuration(durStr string) (time.Duration, error) {
  // Run the duration string through a bunch of processing to get it into a time.Duration format that can be parsed by go
  durStr = strings.Replace(durStr, " seconds", "s", -1)
  durStr = strings.Replace(durStr, " second", "s", -1)
  durStr = strings.Replace(durStr, " minutes", "m", -1)
  durStr = strings.Replace(durStr, " minute", "m", -1)
  durStr = strings.Replace(durStr, " hours", "h", -1)
  durStr = strings.Replace(durStr, " hour", "h", -1)
  durStr = strings.Replace(durStr, " ", "", -1)
  d, err := time.ParseDuration(durStr)
  if err != nil {
    return d, errors.New("Apologies, but I can't seem to parse your duration string.")
  }
  if d.Hours() < 0 || d.Minutes() < 0 || d.Seconds() < 0 {
    return d, errors.New("Apologies, but my functionality does not include the recognition of negative time.")
  }
  return d, nil
}

// This is a horribly complex function to convert an "absolute time" into a go time
// It uses a combination of jinzhu's NOW library and some custom parsing code to
// give the best user experience possible.
func StringToTime(ts string) (time.Time, error) {
  defaultErr := errors.New("Apologies, but I can't seem to read the time you gave me.")

  // Case 1: "Remind me at 8 to do X"
  // NOW will parse this to always mean "8AM of the current day", but I want this to actually mean
  //  - "8PM Today if it is after 8am today but before 8pm today"
  //  - "8AM Tomorrow if it is after 8AM today and after 8pm today"
  works, t, err := parseAbsTimeLoneNumber(ts, defaultErr)
  if err != nil {
    return t, err
  }
  if works {
    return t, nil
  }

  // At this point we pass the timestamp over to NOW to parse
  t, err = now.Parse(ts)
  fmt.Printf("%v\n", t.String())
  if err != nil {
    return t, defaultErr
  }
  return t, nil

}

func parseAbsTimeLoneNumber(ts string, defaultErr error) (bool, time.Time, error) {
  if NewRegex("^[0-9]{1,2}$").Matches(ts) || NewRegex("^[0-9]{1,2}:[0-9]{2}$").Matches(ts) {
    log.Trace("Parsing absolute time assuming lone number")
    t, err := now.Parse(ts)
    if err != nil {
      log.Trace("Error: %v\n", err.Error())
      return false, t, defaultErr
    }
    if t.Before(time.Now()) {
      t = t.Add(12 * time.Hour)
    }
    if t.Before(time.Now()) {
      t = t.Add(12 * time.Hour)
    }
    return true, t, nil
  } else {
    return false, time.Now(), nil
  }
}
<|endoftext|>


<commit_before>package glacier

import (
	"crypto/sha256"
	"encoding/hex"
	"hash"
)

// MultiTreeHasher is used to calculate tree hashes for multi-part uploads
// Call Add sequentially on hashes you have calculated them for
// parts individually, and CreateHash to get the resulting root-level
// hash to use in a CompleteMultipart request.
type MultiTreeHasher struct {
	nodes [][sha256.Size]byte
}

// Add appends the hex-encoded hash to the treehash as a new node
// Add must be called sequentially on parts.
func (t *MultiTreeHasher) Add(hash string) {
	var b [sha256.Size]byte
	hex.Decode(b[:], []byte(hash))
	t.nodes = append(t.nodes, b)
}

// CreateHash returns the root-level hex-encoded hash to send in the
// CompleteMultipart request.
func (t *MultiTreeHasher) CreateHash() string {
	if len(t.nodes) == 0 {
		return ""
	}
	rootHash := treeHash(t.nodes)
	return hex.EncodeToString(rootHash[:])
}

// treeHash calculates the root-level treeHash given sequential
// leaf nodes.
func treeHash(nodes [][sha256.Size]byte) [sha256.Size]byte {
	var combine [sha256.Size * 2]byte
	for len(nodes) > 1 {
		for i := 0; i < len(nodes)/2; i++ {
			copy(combine[:sha256.Size], nodes[i*2][:])
			copy(combine[sha256.Size:], nodes[i*2+1][:])
			nodes[i] = sha256.Sum256(combine[:])
		}
		if len(nodes)%2 == 0 {
			nodes = nodes[:len(nodes)/2]
		} else {
			nodes[len(nodes)/2] = nodes[len(nodes)-1]
			nodes = nodes[:len(nodes)/2+1]
		}
	}
	return nodes[0]
}

// TreeHash is used to calculate the tree hash and regular sha256 hash of the
// data written to it. These values are needed when uploading an archive or
// verifying an aligned download. First each 1 MiB chunk of data is hashed.
// Second each consecutive child node's hashes are concatenated then hashed (if
// there is a single node left it is promoted to the next level). The second
// step is repeated until there is only a single node, this is the tree hash.
// See docs.aws.amazon.com/amazonglacier/latest/dev/checksum-calculations.html
type TreeHash struct {
	nodes       [][sha256.Size]byte
	remaining   []byte
	runningHash hash.Hash         // linear
	treeHash    [sha256.Size]byte // computed
	linearHash  []byte            //computed
}

// NewTreeHash returns an new, initialized tree hasher.
func NewTreeHash() *TreeHash {
	result := &TreeHash{}
	result.Reset()
	return result
}

// Reset the tree hash's state allowing it to be reused.
func (th *TreeHash) Reset() {
	th.runningHash = sha256.New()
	th.remaining = make([]byte, 0)
	th.nodes = make([][sha256.Size]byte, 0)
	th.treeHash = [sha256.Size]byte{}
	th.linearHash = make([]byte, 0)
}

// Write writes all of p, storing every 1 MiB of data's hash.
func (th *TreeHash) Write(p []byte) (n int, err error) {
	n = len(p)
	th.remaining = append(th.remaining, p...)

	// Append one-megabyte increments to the hashes.
	for len(th.remaining) >= (1 << 20) {
		th.nodes = append(th.nodes, sha256.Sum256(th.remaining[:1<<20]))
		th.runningHash.Write(th.remaining[:1<<20])
		th.remaining = th.remaining[1<<20:]
	}
	return
}

// Close closes the the remaing chunks of data and then calculates the tree hash.
func (th *TreeHash) Close() error {
	// create last node; it is impossible that it has a size > 1 MB
	if len(th.remaining) > 0 {
		th.nodes = append(th.nodes, sha256.Sum256(th.remaining))
		th.runningHash.Write(th.remaining)
		th.remaining = make([]byte, 0)
	}
	// Calculate the tree and linear hashes
	if len(th.nodes) > 0 {
		th.treeHash = treeHash(th.nodes)
	}
	th.linearHash = th.runningHash.Sum(nil)
	return nil
}

// TreeHash returns the root-level tree hash of everything written.
func (th *TreeHash) TreeHash() []byte {
	return th.treeHash[:]
}

// Hash returns the linear sha256 checksum of everything written.
func (th *TreeHash) Hash() []byte {
	return th.linearHash[:]
}
<commit_msg>Reduce allocations when writing to and reseting tree hash.<commit_after>package glacier

import (
	"crypto/sha256"
	"encoding/hex"
	"hash"
)

// MultiTreeHasher is used to calculate tree hashes for multi-part uploads
// Call Add sequentially on hashes you have calculated them for
// parts individually, and CreateHash to get the resulting root-level
// hash to use in a CompleteMultipart request.
type MultiTreeHasher struct {
	nodes [][sha256.Size]byte
}

// Add appends the hex-encoded hash to the treehash as a new node
// Add must be called sequentially on parts.
func (t *MultiTreeHasher) Add(hash string) {
	var b [sha256.Size]byte
	hex.Decode(b[:], []byte(hash))
	t.nodes = append(t.nodes, b)
}

// CreateHash returns the root-level hex-encoded hash to send in the
// CompleteMultipart request.
func (t *MultiTreeHasher) CreateHash() string {
	if len(t.nodes) == 0 {
		return ""
	}
	rootHash := treeHash(t.nodes)
	return hex.EncodeToString(rootHash[:])
}

// treeHash calculates the root-level treeHash given sequential
// leaf nodes.
func treeHash(nodes [][sha256.Size]byte) [sha256.Size]byte {
	var combine [sha256.Size * 2]byte
	for len(nodes) > 1 {
		for i := 0; i < len(nodes)/2; i++ {
			copy(combine[:sha256.Size], nodes[i*2][:])
			copy(combine[sha256.Size:], nodes[i*2+1][:])
			nodes[i] = sha256.Sum256(combine[:])
		}
		if len(nodes)%2 == 0 {
			nodes = nodes[:len(nodes)/2]
		} else {
			nodes[len(nodes)/2] = nodes[len(nodes)-1]
			nodes = nodes[:len(nodes)/2+1]
		}
	}
	return nodes[0]
}

// TreeHash is used to calculate the tree hash and regular sha256 hash of the
// data written to it. These values are needed when uploading an archive or
// verifying an aligned download. First each 1 MiB chunk of data is hashed.
// Second each consecutive child node's hashes are concatenated then hashed (if
// there is a single node left it is promoted to the next level). The second
// step is repeated until there is only a single node, this is the tree hash.
// See docs.aws.amazon.com/amazonglacier/latest/dev/checksum-calculations.html
type TreeHash struct {
	remaining   []byte
	nodes       [][sha256.Size]byte
	runningHash hash.Hash         // linear
	treeHash    [sha256.Size]byte // computed
	linearHash  [sha256.Size]byte // computed
}

// NewTreeHash returns an new, initialized tree hasher.
func NewTreeHash() *TreeHash {
	result := &TreeHash{
		runningHash: sha256.New(),
		remaining:   make([]byte, 0, 1<<20),
	}
	result.Reset()
	return result
}

// Reset the tree hash's state allowing it to be reused.
func (th *TreeHash) Reset() {
	th.runningHash.Reset()
	th.remaining = th.remaining[:0]
	th.nodes = th.nodes[:0]
	th.treeHash = [sha256.Size]byte{}
	th.linearHash = [sha256.Size]byte{}
}

// Write writes all of p, storing every 1 MiB of data's hash.
func (th *TreeHash) Write2(p []byte) (n int, err error) {
	n = len(p)
	th.remaining = append(th.remaining, p...)

	// Append one-megabyte increments to the hashes.
	for len(th.remaining) >= (1 << 20) {
		th.nodes = append(th.nodes, sha256.Sum256(th.remaining[:1<<20]))
		th.runningHash.Write(th.remaining[:1<<20])
		th.remaining = th.remaining[1<<20:]
	}
	return
}

// Write writes all of p, storing every 1 MiB of data's hash.
func (th *TreeHash) Write(p []byte) (int, error) {
	n := len(p)

	// Not enough data to fill a 1 MB chunk.
	if len(th.remaining)+len(p) < 1<<20 {
		th.remaining = append(th.remaining, p...)
		return n, nil
	}

	// Move enough to fill th.remaining to 1 MB.
	fill := 1<<20 - len(th.remaining)
	th.remaining = append(th.remaining, p[:fill]...)
	p = p[fill:]

	// Append the 1 MB in th.remaining.
	th.nodes = append(th.nodes, sha256.Sum256(th.remaining))
	th.runningHash.Write(th.remaining)
	th.remaining = th.remaining[:0]

	// Append all 1M chunks remaining in p.
	for len(p) >= 1<<20 {
		th.nodes = append(th.nodes, sha256.Sum256(p[:1<<20]))
		th.runningHash.Write(p[:1<<20])
		p = p[1<<20:]
	}

	// Copy what remains in p to th.remaining.
	th.remaining = append(th.remaining, p...)

	return n, nil
}

// Close closes the the remaing chunks of data and then calculates the tree hash.
func (th *TreeHash) Close() error {
	// create last node; it is impossible that it has a size > 1 MB
	if len(th.remaining) > 0 {
		th.nodes = append(th.nodes, sha256.Sum256(th.remaining))
		th.runningHash.Write(th.remaining)
	}
	// Calculate the tree and linear hashes
	if len(th.nodes) > 0 {
		th.treeHash = treeHash(th.nodes)
	}
	th.runningHash.Sum(th.linearHash[:0])
	return nil
}

// TreeHash returns the root-level tree hash of everything written.
func (th *TreeHash) TreeHash() []byte {
	return th.treeHash[:]
}

// Hash returns the linear sha256 checksum of everything written.
func (th *TreeHash) Hash() []byte {
	return th.linearHash[:]
}
<|endoftext|>


<commit_before>package x86

import (
	cs "github.com/bnagy/gapstone"
	uc "github.com/unicorn-engine/unicorn/bindings/go/unicorn"

	"github.com/lunixbochs/usercorn/go/models"
)

var Arch = &models.Arch{
	Bits:    32,
	Radare:  "x86",
	CS_ARCH: cs.CS_ARCH_X86,
	CS_MODE: cs.CS_MODE_32,
	UC_ARCH: uc.ARCH_X86,
	UC_MODE: uc.MODE_32,
	PC:      uc.X86_REG_EIP,
	SP:      uc.X86_REG_ESP,
	Regs: map[string]int{
		"eip": uc.X86_REG_EIP,
		"esp": uc.X86_REG_ESP,
		"eax": uc.X86_REG_EAX,
		"ebx": uc.X86_REG_EBX,
		"ecx": uc.X86_REG_ECX,
		"edx": uc.X86_REG_EDX,
		"esi": uc.X86_REG_ESI,
		"edi": uc.X86_REG_EDI,
	},
	DefaultRegs: []string{
		"eax", "ebx", "ecx", "edx", "esi", "edi",
	},
}
<commit_msg>include ebp in x86 regs<commit_after>package x86

import (
	cs "github.com/bnagy/gapstone"
	uc "github.com/unicorn-engine/unicorn/bindings/go/unicorn"

	"github.com/lunixbochs/usercorn/go/models"
)

var Arch = &models.Arch{
	Bits:    32,
	Radare:  "x86",
	CS_ARCH: cs.CS_ARCH_X86,
	CS_MODE: cs.CS_MODE_32,
	UC_ARCH: uc.ARCH_X86,
	UC_MODE: uc.MODE_32,
	PC:      uc.X86_REG_EIP,
	SP:      uc.X86_REG_ESP,
	Regs: map[string]int{
		"eip": uc.X86_REG_EIP,
		"esp": uc.X86_REG_ESP,
		"ebp": uc.X86_REG_EBP,
		"eax": uc.X86_REG_EAX,
		"ebx": uc.X86_REG_EBX,
		"ecx": uc.X86_REG_ECX,
		"edx": uc.X86_REG_EDX,
		"esi": uc.X86_REG_ESI,
		"edi": uc.X86_REG_EDI,
	},
	DefaultRegs: []string{
		"eax", "ebx", "ecx", "edx", "esi", "edi", "ebp",
	},
}
<|endoftext|>


<commit_before>// Copyright 2016 Keybase, Inc. All rights reserved. Use of
// this source code is governed by the included BSD license.

package badges

import (
	"golang.org/x/net/context"

	"github.com/keybase/client/go/libkb"
	"github.com/keybase/client/go/protocol/chat1"
	"github.com/keybase/client/go/protocol/gregor1"
	"github.com/keybase/client/go/protocol/keybase1"
)

// Badger keeps a BadgeState up to date and broadcasts it to electron.
// This is the client-specific glue.
// The state is kept up to date by subscribing to:
// - All gregor state updates
// - All chat.activity gregor OOBMs
// - Logout
type Badger struct {
	libkb.Contextified
	badgeState *BadgeState
}

func NewBadger(g *libkb.GlobalContext) *Badger {
	return &Badger{
		Contextified: libkb.NewContextified(g),
		badgeState:   NewBadgeState(g.Log),
	}
}

func (b *Badger) PushState(state gregor1.State) {
	b.G().Log.Debug("Badger update with gregor state")
	b.badgeState.UpdateWithGregor(state)
	err := b.Send()
	if err != nil {
		b.G().Log.Warning("Badger send (pushstate) failed: %v", err)
	}
}

func (b *Badger) PushChatUpdate(update chat1.UnreadUpdate, inboxVers chat1.InboxVers) {
	b.G().Log.Debug("Badger update with chat update")
	b.badgeState.UpdateWithChat(update, inboxVers)
	err := b.Send()
	if err != nil {
		b.G().Log.Warning("Badger send (pushchatupdate) failed: %v", err)
	}
}

func (b *Badger) Resync(ctx context.Context, remoteClient *chat1.RemoteClient) error {
	b.G().Log.Debug("Badger resync req")
	update, err := remoteClient.GetUnreadUpdateFull(ctx, chat1.InboxVers(0))
	if err != nil {
		b.G().Log.Warning("Badger resync failed: %v", err)
		return err
	}
	b.badgeState.UpdateWithChatFull(update)
	err = b.Send()
	if err != nil {
		b.G().Log.Warning("Badger send (resync) failed: %v", err)
	} else {
		b.G().Log.Debug("Badger resync complete")
	}
	return err
}

func (b *Badger) Clear(ctx context.Context) {
	b.badgeState.Clear()
	err := b.Send()
	if err != nil {
		b.G().Log.Warning("Badger send (clear) failed: %v", err)
	}
}

// Send the badgestate to electron
func (b *Badger) Send() error {
	state, err := b.badgeState.Export()
	if err != nil {
		return err
	}
	b.log(state)
	b.G().NotifyRouter.HandleBadgeState(state)
	return nil
}

func (b *Badger) State() *BadgeState {
	return b.badgeState
}

// Log a copy of the badgestate with some zeros stripped off for brevity.
func (b *Badger) log(state1 keybase1.BadgeState) {
	var state2 keybase1.BadgeState
	state2 = state1
	state2.Conversations = nil
	for _, c1 := range state1.Conversations {
		if c1.UnreadMessages == 0 {
			continue
		}
		c2id := c1.ConvID
		if len(c1.ConvID) >= chat1.DbShortFormLen {
			// This is the db short form for logging brevity only.
			// Don't let this leave this method.
			c2id = chat1.ConversationID([]byte(c1.ConvID)).DbShortForm()
		}
		c2 := keybase1.BadgeConversationInfo{
			ConvID:         c2id,
			UnreadMessages: c1.UnreadMessages,
		}
		state2.Conversations = append(state2.Conversations, c2)
	}
	b.G().Log.Debug("Badger send: %+v", state2)
}
<commit_msg>send the current inbox version up to GetUnreadUpdateFull CORE-4784 (#6362)<commit_after>// Copyright 2016 Keybase, Inc. All rights reserved. Use of
// this source code is governed by the included BSD license.

package badges

import (
	"golang.org/x/net/context"

	"github.com/keybase/client/go/chat/storage"
	"github.com/keybase/client/go/libkb"
	"github.com/keybase/client/go/protocol/chat1"
	"github.com/keybase/client/go/protocol/gregor1"
	"github.com/keybase/client/go/protocol/keybase1"
)

// Badger keeps a BadgeState up to date and broadcasts it to electron.
// This is the client-specific glue.
// The state is kept up to date by subscribing to:
// - All gregor state updates
// - All chat.activity gregor OOBMs
// - Logout
type Badger struct {
	libkb.Contextified
	badgeState *BadgeState
}

func NewBadger(g *libkb.GlobalContext) *Badger {
	return &Badger{
		Contextified: libkb.NewContextified(g),
		badgeState:   NewBadgeState(g.Log),
	}
}

func (b *Badger) PushState(state gregor1.State) {
	b.G().Log.Debug("Badger update with gregor state")
	b.badgeState.UpdateWithGregor(state)
	err := b.Send()
	if err != nil {
		b.G().Log.Warning("Badger send (pushstate) failed: %v", err)
	}
}

func (b *Badger) PushChatUpdate(update chat1.UnreadUpdate, inboxVers chat1.InboxVers) {
	b.G().Log.Debug("Badger update with chat update")
	b.badgeState.UpdateWithChat(update, inboxVers)
	err := b.Send()
	if err != nil {
		b.G().Log.Warning("Badger send (pushchatupdate) failed: %v", err)
	}
}

func (b *Badger) inboxVersion(ctx context.Context) chat1.InboxVers {
	uid := b.G().Env.GetUID()
	vers, err := storage.NewInbox(b.G(), uid.ToBytes()).Version(ctx)
	if err != nil {
		b.G().Log.Debug("Badger: inboxVersion error: %s", err.Error())
		return chat1.InboxVers(0)
	}
	return vers
}

func (b *Badger) Resync(ctx context.Context, remoteClient *chat1.RemoteClient) error {
	b.G().Log.Debug("Badger resync req")
	iboxVersion := b.inboxVersion(ctx)
	b.G().Log.Debug("Badger: Resync(): using inbox version: %v", iboxVersion)
	update, err := remoteClient.GetUnreadUpdateFull(ctx, iboxVersion)
	if err != nil {
		b.G().Log.Warning("Badger resync failed: %v", err)
		return err
	}
	b.badgeState.UpdateWithChatFull(update)
	err = b.Send()
	if err != nil {
		b.G().Log.Warning("Badger send (resync) failed: %v", err)
	} else {
		b.G().Log.Debug("Badger resync complete")
	}
	return err
}

func (b *Badger) Clear(ctx context.Context) {
	b.badgeState.Clear()
	err := b.Send()
	if err != nil {
		b.G().Log.Warning("Badger send (clear) failed: %v", err)
	}
}

// Send the badgestate to electron
func (b *Badger) Send() error {
	state, err := b.badgeState.Export()
	if err != nil {
		return err
	}
	b.log(state)
	b.G().NotifyRouter.HandleBadgeState(state)
	return nil
}

func (b *Badger) State() *BadgeState {
	return b.badgeState
}

// Log a copy of the badgestate with some zeros stripped off for brevity.
func (b *Badger) log(state1 keybase1.BadgeState) {
	var state2 keybase1.BadgeState
	state2 = state1
	state2.Conversations = nil
	for _, c1 := range state1.Conversations {
		if c1.UnreadMessages == 0 {
			continue
		}
		c2id := c1.ConvID
		if len(c1.ConvID) >= chat1.DbShortFormLen {
			// This is the db short form for logging brevity only.
			// Don't let this leave this method.
			c2id = chat1.ConversationID([]byte(c1.ConvID)).DbShortForm()
		}
		c2 := keybase1.BadgeConversationInfo{
			ConvID:         c2id,
			UnreadMessages: c1.UnreadMessages,
		}
		state2.Conversations = append(state2.Conversations, c2)
	}
	b.G().Log.Debug("Badger send: %+v", state2)
}
<|endoftext|>


<commit_before>// Copyright 2015 Keybase, Inc. All rights reserved. Use of
// this source code is governed by the included BSD license.

package engine

import (
	"fmt"

	"github.com/keybase/client/go/libkb"
	keybase1 "github.com/keybase/client/go/protocol"
	triplesec "github.com/keybase/go-triplesec"
)

type SignupEngine struct {
	pwsalt        []byte
	ppStream      *libkb.PassphraseStream
	tsec          *triplesec.Cipher
	uid           keybase1.UID
	me            *libkb.User
	signingKey    libkb.GenericKey
	encryptionKey libkb.GenericKey
	arg           *SignupEngineRunArg
	lks           *libkb.LKSec
	libkb.Contextified
}

type SignupEngineRunArg struct {
	Username    string
	Email       string
	InviteCode  string
	Passphrase  string
	StoreSecret bool
	DeviceName  string
	SkipGPG     bool
	SkipMail    bool
	SkipPaper   bool
}

func NewSignupEngine(arg *SignupEngineRunArg, g *libkb.GlobalContext) *SignupEngine {
	return &SignupEngine{
		arg:          arg,
		Contextified: libkb.NewContextified(g),
	}
}

func (s *SignupEngine) Name() string {
	return "Signup"
}

func (s *SignupEngine) RequiredUIs() []libkb.UIKind {
	return nil
}

func (s *SignupEngine) Prereqs() Prereqs { return Prereqs{} }

func (s *SignupEngine) SubConsumers() []libkb.UIConsumer {
	return []libkb.UIConsumer{
		&GPGImportKeyEngine{},
		&DeviceWrap{},
		&PaperKeyPrimary{},
	}
}

func (s *SignupEngine) GetMe() *libkb.User {
	return s.me
}

func (s *SignupEngine) Run(ctx *Context) error {
	// make sure we're starting with a clear login state:
	if err := s.G().Logout(); err != nil {
		return err
	}

	f := func(a libkb.LoginContext) error {
		if err := s.genPassphraseStream(a, s.arg.Passphrase); err != nil {
			return err
		}

		if err := s.join(a, s.arg.Username, s.arg.Email, s.arg.InviteCode, s.arg.SkipMail); err != nil {
			return err
		}

		if err := s.registerDevice(a, ctx, s.arg.DeviceName); err != nil {
			return err
		}

		s.G().NotifyRouter.HandleLogin(s.arg.Username)

		if !s.arg.SkipPaper {
			if err := s.genPaperKeys(ctx); err != nil {
				return err
			}
		}

		if s.arg.SkipGPG {
			return nil
		}

		if wantsGPG, err := s.checkGPG(ctx); err != nil {
			return err
		} else if wantsGPG {
			if err := s.addGPG(a, ctx, true); err != nil {
				return fmt.Errorf("addGPG error: %s", err)
			}
		}

		return nil
	}
	return s.G().LoginState().ExternalFunc(f, "SignupEngine - Run")
}

func (s *SignupEngine) genPassphraseStream(a libkb.LoginContext, passphrase string) error {
	salt, err := libkb.RandBytes(triplesec.SaltLen)
	if err != nil {
		return err
	}
	s.pwsalt = salt
	s.tsec, s.ppStream, err = libkb.StretchPassphrase(passphrase, salt)
	if err != nil {
		return err
	}
	return nil
}

func (s *SignupEngine) join(a libkb.LoginContext, username, email, inviteCode string, skipMail bool) error {
	joinEngine := NewSignupJoinEngine(s.G())

	arg := SignupJoinEngineRunArg{
		Username:   username,
		Email:      email,
		InviteCode: inviteCode,
		PWHash:     s.ppStream.PWHash(),
		PWSalt:     s.pwsalt,
		SkipMail:   skipMail,
	}
	res := joinEngine.Run(a, arg)
	if res.Err != nil {
		return res
	}

	s.ppStream.SetGeneration(res.PpGen)
	a.CreateStreamCache(s.tsec, s.ppStream)

	s.uid = res.UID
	s.G().Log.Debug("contextified: %v\n", s.G())
	user, err := libkb.LoadUser(libkb.LoadUserArg{Self: true, UID: res.UID, PublicKeyOptional: true, Contextified: libkb.NewContextified(s.G())})
	if err != nil {
		return err
	}
	s.me = user
	return nil
}

func (s *SignupEngine) registerDevice(a libkb.LoginContext, ctx *Context, deviceName string) error {
	s.lks = libkb.NewLKSec(s.ppStream, s.uid, s.G())
	args := &DeviceWrapArgs{
		Me:         s.me,
		DeviceName: deviceName,
		DeviceType: libkb.DeviceTypeDesktop,
		Lks:        s.lks,
		IsEldest:   true,
	}
	eng := NewDeviceWrap(args, s.G())
	ctx.LoginContext = a
	if err := RunEngine(eng, ctx); err != nil {
		return err
	}
	s.signingKey = eng.SigningKey()
	s.encryptionKey = eng.EncryptionKey()

	if err := ctx.LoginContext.LocalSession().SetDeviceProvisioned(s.G().Env.GetDeviceID()); err != nil {
		// this isn't a fatal error, session will stay in memory...
		s.G().Log.Warning("error saving session file: %s", err)
	}

	if s.arg.StoreSecret {
		// Create the secret store as late as possible here
		// (instead of when we first get the value of
		// StoreSecret) as the username may change during the
		// signup process.
		secretStore := libkb.NewSecretStore(s.G(), s.me.GetNormalizedName())
		secret, err := s.lks.GetSecret(a)
		if err != nil {
			return err
		}
		// Ignore any errors storing the secret.
		storeSecretErr := secretStore.StoreSecret(secret)
		if storeSecretErr != nil {
			s.G().Log.Warning("StoreSecret error: %s", storeSecretErr)
		}
	}

	// is there any reason *not* to do this?
	ctx.LoginContext.SetCachedSecretKey(libkb.SecretKeyArg{KeyType: libkb.DeviceSigningKeyType}, s.signingKey)
	ctx.LoginContext.SetCachedSecretKey(libkb.SecretKeyArg{KeyType: libkb.DeviceEncryptionKeyType}, eng.EncryptionKey())

	s.G().Log.Debug("registered new device: %s", s.G().Env.GetDeviceID())
	s.G().Log.Debug("eldest kid: %s", s.me.GetEldestKID())

	return nil
}

func (s *SignupEngine) genPaperKeys(ctx *Context) error {
	args := &PaperKeyPrimaryArgs{
		Me:         s.me,
		SigningKey: s.signingKey,
	}
	eng := NewPaperKeyPrimary(s.G(), args)
	return RunEngine(eng, ctx)
}

func (s *SignupEngine) checkGPG(ctx *Context) (bool, error) {
	eng := NewGPGImportKeyEngine(nil, s.G())
	return eng.WantsGPG(ctx)
}

func (s *SignupEngine) addGPG(lctx libkb.LoginContext, ctx *Context, allowMulti bool) error {
	s.G().Log.Debug("SignupEngine.addGPG.  signingKey: %v\n", s.signingKey)
	arg := GPGImportKeyArg{Signer: s.signingKey, AllowMulti: allowMulti, Me: s.me, Lks: s.lks}
	eng := NewGPGImportKeyEngine(&arg, s.G())
	ctx.LoginContext = lctx
	if err := RunEngine(eng, ctx); err != nil {
		return err
	}

	if s.signingKey == nil {
		s.signingKey = eng.LastKey()
	}
	return nil
}
<commit_msg>Move login notification to a safe place<commit_after>// Copyright 2015 Keybase, Inc. All rights reserved. Use of
// this source code is governed by the included BSD license.

package engine

import (
	"fmt"

	"github.com/keybase/client/go/libkb"
	keybase1 "github.com/keybase/client/go/protocol"
	triplesec "github.com/keybase/go-triplesec"
)

type SignupEngine struct {
	pwsalt        []byte
	ppStream      *libkb.PassphraseStream
	tsec          *triplesec.Cipher
	uid           keybase1.UID
	me            *libkb.User
	signingKey    libkb.GenericKey
	encryptionKey libkb.GenericKey
	arg           *SignupEngineRunArg
	lks           *libkb.LKSec
	libkb.Contextified
}

type SignupEngineRunArg struct {
	Username    string
	Email       string
	InviteCode  string
	Passphrase  string
	StoreSecret bool
	DeviceName  string
	SkipGPG     bool
	SkipMail    bool
	SkipPaper   bool
}

func NewSignupEngine(arg *SignupEngineRunArg, g *libkb.GlobalContext) *SignupEngine {
	return &SignupEngine{
		arg:          arg,
		Contextified: libkb.NewContextified(g),
	}
}

func (s *SignupEngine) Name() string {
	return "Signup"
}

func (s *SignupEngine) RequiredUIs() []libkb.UIKind {
	return nil
}

func (s *SignupEngine) Prereqs() Prereqs { return Prereqs{} }

func (s *SignupEngine) SubConsumers() []libkb.UIConsumer {
	return []libkb.UIConsumer{
		&GPGImportKeyEngine{},
		&DeviceWrap{},
		&PaperKeyPrimary{},
	}
}

func (s *SignupEngine) GetMe() *libkb.User {
	return s.me
}

func (s *SignupEngine) Run(ctx *Context) error {
	// make sure we're starting with a clear login state:
	if err := s.G().Logout(); err != nil {
		return err
	}

	f := func(a libkb.LoginContext) error {
		if err := s.genPassphraseStream(a, s.arg.Passphrase); err != nil {
			return err
		}

		if err := s.join(a, s.arg.Username, s.arg.Email, s.arg.InviteCode, s.arg.SkipMail); err != nil {
			return err
		}

		if err := s.registerDevice(a, ctx, s.arg.DeviceName); err != nil {
			return err
		}

		if !s.arg.SkipPaper {
			if err := s.genPaperKeys(ctx); err != nil {
				return err
			}
		}

		if s.arg.SkipGPG {
			return nil
		}

		if wantsGPG, err := s.checkGPG(ctx); err != nil {
			return err
		} else if wantsGPG {
			if err := s.addGPG(a, ctx, true); err != nil {
				return fmt.Errorf("addGPG error: %s", err)
			}
		}

		return nil
	}

	if err := s.G().LoginState().ExternalFunc(f, "SignupEngine - Run"); err != nil {
		return err
	}

	// signup complete, notify anyone interested.
	// (and don't notify inside a LoginState action to avoid
	// a chance of timing out)
	s.G().NotifyRouter.HandleLogin(s.arg.Username)

	return nil

}

func (s *SignupEngine) genPassphraseStream(a libkb.LoginContext, passphrase string) error {
	salt, err := libkb.RandBytes(triplesec.SaltLen)
	if err != nil {
		return err
	}
	s.pwsalt = salt
	s.tsec, s.ppStream, err = libkb.StretchPassphrase(passphrase, salt)
	if err != nil {
		return err
	}
	return nil
}

func (s *SignupEngine) join(a libkb.LoginContext, username, email, inviteCode string, skipMail bool) error {
	joinEngine := NewSignupJoinEngine(s.G())

	arg := SignupJoinEngineRunArg{
		Username:   username,
		Email:      email,
		InviteCode: inviteCode,
		PWHash:     s.ppStream.PWHash(),
		PWSalt:     s.pwsalt,
		SkipMail:   skipMail,
	}
	res := joinEngine.Run(a, arg)
	if res.Err != nil {
		return res
	}

	s.ppStream.SetGeneration(res.PpGen)
	a.CreateStreamCache(s.tsec, s.ppStream)

	s.uid = res.UID
	s.G().Log.Debug("contextified: %v\n", s.G())
	user, err := libkb.LoadUser(libkb.LoadUserArg{Self: true, UID: res.UID, PublicKeyOptional: true, Contextified: libkb.NewContextified(s.G())})
	if err != nil {
		return err
	}
	s.me = user
	return nil
}

func (s *SignupEngine) registerDevice(a libkb.LoginContext, ctx *Context, deviceName string) error {
	s.lks = libkb.NewLKSec(s.ppStream, s.uid, s.G())
	args := &DeviceWrapArgs{
		Me:         s.me,
		DeviceName: deviceName,
		DeviceType: libkb.DeviceTypeDesktop,
		Lks:        s.lks,
		IsEldest:   true,
	}
	eng := NewDeviceWrap(args, s.G())
	ctx.LoginContext = a
	if err := RunEngine(eng, ctx); err != nil {
		return err
	}
	s.signingKey = eng.SigningKey()
	s.encryptionKey = eng.EncryptionKey()

	if err := ctx.LoginContext.LocalSession().SetDeviceProvisioned(s.G().Env.GetDeviceID()); err != nil {
		// this isn't a fatal error, session will stay in memory...
		s.G().Log.Warning("error saving session file: %s", err)
	}

	if s.arg.StoreSecret {
		// Create the secret store as late as possible here
		// (instead of when we first get the value of
		// StoreSecret) as the username may change during the
		// signup process.
		secretStore := libkb.NewSecretStore(s.G(), s.me.GetNormalizedName())
		secret, err := s.lks.GetSecret(a)
		if err != nil {
			return err
		}
		// Ignore any errors storing the secret.
		storeSecretErr := secretStore.StoreSecret(secret)
		if storeSecretErr != nil {
			s.G().Log.Warning("StoreSecret error: %s", storeSecretErr)
		}
	}

	// is there any reason *not* to do this?
	ctx.LoginContext.SetCachedSecretKey(libkb.SecretKeyArg{KeyType: libkb.DeviceSigningKeyType}, s.signingKey)
	ctx.LoginContext.SetCachedSecretKey(libkb.SecretKeyArg{KeyType: libkb.DeviceEncryptionKeyType}, eng.EncryptionKey())

	s.G().Log.Debug("registered new device: %s", s.G().Env.GetDeviceID())
	s.G().Log.Debug("eldest kid: %s", s.me.GetEldestKID())

	return nil
}

func (s *SignupEngine) genPaperKeys(ctx *Context) error {
	args := &PaperKeyPrimaryArgs{
		Me:         s.me,
		SigningKey: s.signingKey,
	}
	eng := NewPaperKeyPrimary(s.G(), args)
	return RunEngine(eng, ctx)
}

func (s *SignupEngine) checkGPG(ctx *Context) (bool, error) {
	eng := NewGPGImportKeyEngine(nil, s.G())
	return eng.WantsGPG(ctx)
}

func (s *SignupEngine) addGPG(lctx libkb.LoginContext, ctx *Context, allowMulti bool) error {
	s.G().Log.Debug("SignupEngine.addGPG.  signingKey: %v\n", s.signingKey)
	arg := GPGImportKeyArg{Signer: s.signingKey, AllowMulti: allowMulti, Me: s.me, Lks: s.lks}
	eng := NewGPGImportKeyEngine(&arg, s.G())
	ctx.LoginContext = lctx
	if err := RunEngine(eng, ctx); err != nil {
		return err
	}

	if s.signingKey == nil {
		s.signingKey = eng.LastKey()
	}
	return nil
}
<|endoftext|>