<commit_before>package main

import (
	"bufio"
	"encoding/json"
	"flag"
	"fmt"
	"log"
	"net/http"
	"os"
	"strings"
	"sync"
	"time"

	"github.com/mercari/gaurun/gaurun"
	"github.com/mercari/gcm"
)

func pushNotification(wg *sync.WaitGroup, req gaurun.RequestGaurunNotification, logPush gaurun.LogPushEntry, apnsClient *http.Client) {
	var result bool
	switch logPush.Platform {
	case "ios":
		result = pushNotificationIos(apnsClient, req)
	case "android":
		result = pushNotificationAndroid(req)
	}
	if !result {
		msg := fmt.Sprintf("failed to push notification: %s %s %s", logPush.Token, logPush.Platform, logPush.Message)
		log.Println(msg)
	} else {
		msg := fmt.Sprintf("succeeded push notification: %s %s %s", logPush.Token, logPush.Platform, logPush.Message)
		log.Println(msg)
	}

	wg.Done()
}

func pushNotificationAndroid(req gaurun.RequestGaurunNotification) bool {
	data := map[string]interface{}{"message": req.Message}
	msg := gcm.NewMessage(data, req.Tokens...)
	msg.CollapseKey = req.CollapseKey
	msg.DelayWhileIdle = req.DelayWhileIdle
	msg.TimeToLive = req.TimeToLive

	sender := &gcm.Sender{ApiKey: gaurun.ConfGaurun.Android.ApiKey}
	sender.Http = new(http.Client)
	sender.Http.Timeout = time.Duration(gaurun.ConfGaurun.Android.Timeout) * time.Second

	resp, err := sender.SendNoRetry(msg)
	if err != nil {
		return false
	}

	if resp.Failure > 0 {
		return true
	}

	return true
}

func pushNotificationIos(client *http.Client, req gaurun.RequestGaurunNotification) bool {

	service := gaurun.NewApnsServiceHttp2(client)

	for _, token := range req.Tokens {

		headers := gaurun.NewApnsHeadersHttp2(&req)
		payload := gaurun.NewApnsPayloadHttp2(&req)

		err := gaurun.ApnsPushHttp2(token, service, headers, payload)
		if err != nil {
			return false
		}
	}

	return true
}

func main() {
	versionPrinted := flag.Bool("v", false, "gaurun version")
	confPath := flag.String("c", "", "configuration file path for gaurun")
	logPath := flag.String("l", "", "log file path for gaurun")
	flag.Parse()

	if *versionPrinted {
		gaurun.PrintVersion()
		os.Exit(0)
	}

	// set default parameters
	gaurun.ConfGaurun = gaurun.BuildDefaultConf()

	// load configuration
	conf, err := gaurun.LoadConf(gaurun.ConfGaurun, *confPath)
	if err != nil {
		log.Fatal(err)
	}
	gaurun.ConfGaurun = conf

	f, err := os.Open(*logPath)
	if err != nil {
		panic(err)
	}
	defer f.Close()

	scanner := bufio.NewScanner(f)
	accepts := make(map[uint64]gaurun.LogPushEntry)
	successes := make(map[uint64]gaurun.LogPushEntry)

	for scanner.Scan() {
		var logPush gaurun.LogPushEntry
		line := scanner.Text()
		idx := strings.Index(line, " ")
		JSONStr := line[idx+1:]
		err := json.Unmarshal([]byte(JSONStr), &logPush)
		if err != nil {
			log.Printf("JSON parse error(%s)", JSONStr)
			continue
		}
		if logPush.Type == "accepted-request" {
			continue
		}

		switch logPush.Type {
		case "accepted-push":
			accepts[logPush.ID] = logPush
		case "succeeded-push":
			successes[logPush.ID] = logPush
		}
	}

	losts := make(map[uint64]gaurun.LogPushEntry)
	for id, logPush := range accepts {
		if _, ok := successes[id]; !ok {
			losts[id] = logPush
		}
	}

	apnsClient, err := gaurun.NewApnsClientHttp2(
		gaurun.ConfGaurun.Ios.PemCertPath,
		gaurun.ConfGaurun.Ios.PemKeyPath,
	)
	if err != nil {
		log.Fatal(err)
	}
	apnsClient.Timeout = time.Duration(gaurun.ConfGaurun.Ios.Timeout) * time.Second

	wg := new(sync.WaitGroup)
	for _, logPush := range losts {
		tokens := make([]string, 1)
		var platform int
		tokens = append(tokens, logPush.Token)
		switch logPush.Platform {
		case "ios":
			platform = 1
		case "android":
			platform = 2

		}

		req := gaurun.RequestGaurunNotification{
			Tokens:           tokens,
			Platform:         platform,
			Message:          logPush.Message,
			CollapseKey:      logPush.CollapseKey,
			DelayWhileIdle:   logPush.DelayWhileIdle,
			TimeToLive:       logPush.TimeToLive,
			Badge:            logPush.Badge,
			Sound:            logPush.Sound,
			ContentAvailable: logPush.ContentAvailable,
			Expiry:           logPush.Expiry,
		}
		wg.Add(1)
		go pushNotification(wg, req, logPush, apnsClient)
	}

	wg.Wait()
}
<commit_msg>bugfix: not reached notification with gaurun_recover. refs #37<commit_after>package main

import (
	"bufio"
	"encoding/json"
	"flag"
	"fmt"
	"log"
	"net/http"
	"os"
	"strings"
	"sync"
	"time"

	"github.com/mercari/gaurun/gaurun"
	"github.com/mercari/gcm"
)

func pushNotification(wg *sync.WaitGroup, req gaurun.RequestGaurunNotification, logPush gaurun.LogPushEntry, apnsClient *http.Client) {
	var result bool
	switch logPush.Platform {
	case "ios":
		result = pushNotificationIos(apnsClient, req)
	case "android":
		result = pushNotificationAndroid(req)
	}
	if !result {
		msg := fmt.Sprintf("failed to push notification: %s %s %s", logPush.Token, logPush.Platform, logPush.Message)
		log.Println(msg)
	} else {
		msg := fmt.Sprintf("succeeded push notification: %s %s %s", logPush.Token, logPush.Platform, logPush.Message)
		log.Println(msg)
	}

	wg.Done()
}

func pushNotificationAndroid(req gaurun.RequestGaurunNotification) bool {
	data := map[string]interface{}{"message": req.Message}
	msg := gcm.NewMessage(data, req.Tokens...)
	msg.CollapseKey = req.CollapseKey
	msg.DelayWhileIdle = req.DelayWhileIdle
	msg.TimeToLive = req.TimeToLive

	sender := &gcm.Sender{ApiKey: gaurun.ConfGaurun.Android.ApiKey}
	sender.Http = new(http.Client)
	sender.Http.Timeout = time.Duration(gaurun.ConfGaurun.Android.Timeout) * time.Second

	resp, err := sender.SendNoRetry(msg)
	if err != nil {
		return false
	}

	if resp.Failure > 0 {
		return true
	}

	return true
}

func pushNotificationIos(client *http.Client, req gaurun.RequestGaurunNotification) bool {

	service := gaurun.NewApnsServiceHttp2(client)

	for _, token := range req.Tokens {

		headers := gaurun.NewApnsHeadersHttp2(&req)
		payload := gaurun.NewApnsPayloadHttp2(&req)

		err := gaurun.ApnsPushHttp2(token, service, headers, payload)
		if err != nil {
			return false
		}
	}

	return true
}

func main() {
	versionPrinted := flag.Bool("v", false, "gaurun version")
	confPath := flag.String("c", "", "configuration file path for gaurun")
	logPath := flag.String("l", "", "log file path for gaurun")
	flag.Parse()

	if *versionPrinted {
		gaurun.PrintVersion()
		os.Exit(0)
	}

	// set default parameters
	gaurun.ConfGaurun = gaurun.BuildDefaultConf()

	// load configuration
	conf, err := gaurun.LoadConf(gaurun.ConfGaurun, *confPath)
	if err != nil {
		log.Fatal(err)
	}
	gaurun.ConfGaurun = conf

	f, err := os.Open(*logPath)
	if err != nil {
		panic(err)
	}
	defer f.Close()

	scanner := bufio.NewScanner(f)
	accepts := make(map[uint64]gaurun.LogPushEntry)
	successes := make(map[uint64]gaurun.LogPushEntry)

	for scanner.Scan() {
		var logPush gaurun.LogPushEntry
		line := scanner.Text()
		idx := strings.Index(line, " ")
		JSONStr := line[idx+1:]
		err := json.Unmarshal([]byte(JSONStr), &logPush)
		if err != nil {
			log.Printf("JSON parse error(%s)", JSONStr)
			continue
		}
		if logPush.Type == "accepted-request" {
			continue
		}

		switch logPush.Type {
		case "accepted-push":
			accepts[logPush.ID] = logPush
		case "succeeded-push":
			successes[logPush.ID] = logPush
		}
	}

	losts := make(map[uint64]gaurun.LogPushEntry)
	for id, logPush := range accepts {
		if _, ok := successes[id]; !ok {
			losts[id] = logPush
		}
	}

	apnsClient, err := gaurun.NewApnsClientHttp2(
		gaurun.ConfGaurun.Ios.PemCertPath,
		gaurun.ConfGaurun.Ios.PemKeyPath,
	)
	if err != nil {
		log.Fatal(err)
	}
	apnsClient.Timeout = time.Duration(gaurun.ConfGaurun.Ios.Timeout) * time.Second

	wg := new(sync.WaitGroup)
	for _, logPush := range losts {
		tokens := make([]string, 1)
		var platform int
		tokens[0] = logPush.Token
		switch logPush.Platform {
		case "ios":
			platform = 1
		case "android":
			platform = 2

		}

		req := gaurun.RequestGaurunNotification{
			Tokens:           tokens,
			Platform:         platform,
			Message:          logPush.Message,
			CollapseKey:      logPush.CollapseKey,
			DelayWhileIdle:   logPush.DelayWhileIdle,
			TimeToLive:       logPush.TimeToLive,
			Badge:            logPush.Badge,
			Sound:            logPush.Sound,
			ContentAvailable: logPush.ContentAvailable,
			Expiry:           logPush.Expiry,
		}
		wg.Add(1)
		go pushNotification(wg, req, logPush, apnsClient)
	}

	wg.Wait()
}
<|endoftext|>


<commit_before>package main

import (
	"bufio"
	"encoding/json"
	"flag"
	"fmt"
	"io/ioutil"
	"os"
	"path"
	"strconv"
	"strings"
	"sync"
	"time"

	//"runtime/pprof"

	"github.com/livepeer/go-livepeer/common"
	"github.com/livepeer/lpms/ffmpeg"
	"github.com/livepeer/m3u8"
)

func main() {
	/*
		cprof, err := os.Create("bench.prof")
		if err != nil {
			panic(err)
		}
		pprof.StartCPUProfile(cprof)
		defer pprof.StopCPUProfile()
	*/
	// Override the default flag set since there are dependencies that
	// incorrectly add their own flags (specifically, due to the 'testing'
	// package being linked)
	flag.Set("logtostderr", "true")
	flag.CommandLine = flag.NewFlagSet(os.Args[0], flag.ExitOnError)

	fname := flag.String("fname", "", "Input m3u8 manifest file")
	conc := flag.Int("conc", 1, "# of concurrent transcode sessions")
	segs := flag.Int("segs", 0, "Maximum # of segments to transcode (default all)")
	profs := flag.String("profs", "P240p30fps16x9,P360p30fps16x9,P720p30fps16x9", "Transcoding options for broadcast job, or path to json config")
	nvidia := flag.String("nvidia", "", "Comma-separated list of Nvidia GPU device IDs to use for transcoding")

	flag.Parse()

	if *fname == "" {
		panic("Please provide the input manifest as `-fname <input.m3u8>`. See -h or -help for more.")
	}

	profiles := parseVideoProfiles(*profs)

	f, err := os.Open(*fname)
	if err != nil {
		panic(err)
	}
	p, _, err := m3u8.DecodeFrom(bufio.NewReader(f), true)
	if err != nil {
		panic(err)
	}
	pl, ok := p.(*m3u8.MediaPlaylist)
	if !ok {
		panic("Expecting media PL")
	}

	accel := ffmpeg.Software
	devices := []string{}
	if *nvidia != "" {
		accel = ffmpeg.Nvidia
		devices = strings.Split(*nvidia, ",")
	}

	ffmpeg.InitFFmpeg()
	segCount := 0
	var wg sync.WaitGroup
	dir := path.Dir(*fname)
	start := time.Now()
	fmt.Fprintf(os.Stderr, "Program %s Source %s Concurrency %d Profiles %s\n", os.Args[0], *fname, *conc, *profs)
	fmt.Println("time,stream,segment,length")
	for i := 0; i < *conc; i++ {
		wg.Add(1)
		go func(k int, wg *sync.WaitGroup) {
			tc := ffmpeg.NewTranscoder()
			for j, v := range pl.Segments {
				if *segs > 0 && j >= *segs {
					break
				}
				if v == nil {
					continue
				}
				u := path.Join(dir, v.URI)
				//u := v.URI
				in := &ffmpeg.TranscodeOptionsIn{
					Fname: u,
					Accel: accel,
				}
				if ffmpeg.Software != accel {
					in.Device = devices[k%len(devices)]
				}
				profs2opts := func(profs []ffmpeg.VideoProfile) []ffmpeg.TranscodeOptions {
					opts := []ffmpeg.TranscodeOptions{}
					//for n, p := range profs {
					for _, p := range profs {
						o := ffmpeg.TranscodeOptions{
							//Oname: fmt.Sprintf("%s%s_%s_%d_%d_%d.ts", pfx, accelStr, p.Name, n, k, j),
							Oname:        "-",
							Profile:      p,
							Accel:        accel,
							AudioEncoder: ffmpeg.ComponentOptions{Name: "drop"},
							Muxer:        ffmpeg.ComponentOptions{Name: "null"},
						}
						opts = append(opts, o)
					}
					return opts
				}
				out := profs2opts(profiles)
				t := time.Now()
				_, err := tc.Transcode(in, out)
				end := time.Now()
				fmt.Printf("%s,%d,%d,%0.2v\n", end.Format("2006-01-02 15:04:05.999999999"), k, j, end.Sub(t).Seconds())
				if err != nil {
					panic(err)
				}
				segCount++
			}
			tc.StopTranscoder()
			wg.Done()
		}(i, &wg)
		time.Sleep(300 * time.Millisecond)
	}
	wg.Wait()
	fmt.Fprintf(os.Stderr, "Took %v to transcode %v segments\n",
		time.Now().Sub(start).Seconds(), segCount)
}

func parseVideoProfiles(inp string) []ffmpeg.VideoProfile {
	type profilesJson struct {
		Profiles []struct {
			Name    string `json:"name"`
			Width   int    `json:"width"`
			Height  int    `json:"height"`
			Bitrate int    `json:"bitrate"`
			FPS     uint   `json:"fps"`
			FPSDen  uint   `json:"fpsDen"`
			Profile string `json:"profile"`
			GOP     string `json:"gop"`
		} `json:"profiles"`
	}
	profs := []ffmpeg.VideoProfile{}
	if inp != "" {
		// try opening up json file with profiles
		content, err := ioutil.ReadFile(inp)
		if err == nil && len(content) > 0 {
			// parse json profiles
			resp := &profilesJson{}
			err = json.Unmarshal(content, &resp.Profiles)
			if err != nil {
				panic(err)
			}
			for _, profile := range resp.Profiles {
				name := profile.Name
				if name == "" {
					name = "custom_" + common.DefaultProfileName(
						profile.Width,
						profile.Height,
						profile.Bitrate)
				}
				var gop time.Duration
				if profile.GOP != "" {
					if profile.GOP == "intra" {
						gop = ffmpeg.GOPIntraOnly
					} else {
						gopFloat, err := strconv.ParseFloat(profile.GOP, 64)
						if err != nil {
							panic(err)
						}
						if gopFloat <= 0.0 {
							panic("invalid gop value")
						}
						gop = time.Duration(gopFloat * float64(time.Second))
					}
				}
				encodingProfile, err := common.EncoderProfileNameToValue(profile.Profile)
				if err != nil {
					panic(err)
				}
				prof := ffmpeg.VideoProfile{
					Name:         name,
					Bitrate:      fmt.Sprint(profile.Bitrate),
					Framerate:    profile.FPS,
					FramerateDen: profile.FPSDen,
					Resolution:   fmt.Sprintf("%dx%d", profile.Width, profile.Height),
					Profile:      encodingProfile,
					GOP:          gop,
				}
				profs = append(profs, prof)
			}
		} else {
			// check the built-in profiles
			profs = make([]ffmpeg.VideoProfile, 0)
			presets := strings.Split(inp, ",")
			for _, v := range presets {
				if p, ok := ffmpeg.VideoProfileLookup[strings.TrimSpace(v)]; ok {
					profs = append(profs, p)
				}
			}
		}
		if len(profs) <= 0 {
			panic(fmt.Errorf("No transcoding profiles found"))
		}
	}
	return profs
}
<commit_msg>cmd/livepeer_bench: add segment output support<commit_after>package main

import (
	"bufio"
	"encoding/json"
	"flag"
	"fmt"
	"io/ioutil"
	"os"
	"path"
	"strconv"
	"strings"
	"sync"
	"time"

	//"runtime/pprof"

	"github.com/livepeer/go-livepeer/common"
	"github.com/livepeer/lpms/ffmpeg"
	"github.com/livepeer/m3u8"
)

func main() {
	/*
		cprof, err := os.Create("bench.prof")
		if err != nil {
			panic(err)
		}
		pprof.StartCPUProfile(cprof)
		defer pprof.StopCPUProfile()
	*/
	// Override the default flag set since there are dependencies that
	// incorrectly add their own flags (specifically, due to the 'testing'
	// package being linked)
	flag.Set("logtostderr", "true")
	flag.CommandLine = flag.NewFlagSet(os.Args[0], flag.ExitOnError)

	fname := flag.String("fname", "", "Input m3u8 manifest file")
	conc := flag.Int("conc", 1, "# of concurrent transcode sessions")
	segs := flag.Int("segs", 0, "Maximum # of segments to transcode (default all)")
	profs := flag.String("profs", "P240p30fps16x9,P360p30fps16x9,P720p30fps16x9", "Transcoding options for broadcast job, or path to json config")
	nvidia := flag.String("nvidia", "", "Comma-separated list of Nvidia GPU device IDs to use for transcoding")
	outPrefix := flag.String("outPrefix", "", "Output segments' prefix (no segments are generated by default)")

	flag.Parse()

	if *fname == "" {
		panic("Please provide the input manifest as `-fname <input.m3u8>`. See -h or -help for more.")
	}

	profiles := parseVideoProfiles(*profs)

	f, err := os.Open(*fname)
	if err != nil {
		panic(err)
	}
	p, _, err := m3u8.DecodeFrom(bufio.NewReader(f), true)
	if err != nil {
		panic(err)
	}
	pl, ok := p.(*m3u8.MediaPlaylist)
	if !ok {
		panic("Expecting media PL")
	}

	accel := ffmpeg.Software
	devices := []string{}
	if *nvidia != "" {
		accel = ffmpeg.Nvidia
		devices = strings.Split(*nvidia, ",")
	}

	ffmpeg.InitFFmpeg()
	segCount := 0
	var wg sync.WaitGroup
	dir := path.Dir(*fname)
	start := time.Now()
	fmt.Fprintf(os.Stderr, "Program %s Source %s Concurrency %d Profiles %s\n", os.Args[0], *fname, *conc, *profs)
	fmt.Println("time,stream,segment,length")
	for i := 0; i < *conc; i++ {
		wg.Add(1)
		go func(k int, wg *sync.WaitGroup) {
			tc := ffmpeg.NewTranscoder()
			for j, v := range pl.Segments {
				if *segs > 0 && j >= *segs {
					break
				}
				if v == nil {
					continue
				}
				u := path.Join(dir, v.URI)
				in := &ffmpeg.TranscodeOptionsIn{
					Fname: u,
					Accel: accel,
				}
				if ffmpeg.Software != accel {
					in.Device = devices[k%len(devices)]
				}
				profs2opts := func(profs []ffmpeg.VideoProfile) []ffmpeg.TranscodeOptions {
					opts := []ffmpeg.TranscodeOptions{}
					for n, p := range profs {
						oname := ""
						muxer := ""
						if *outPrefix != "" {
							oname = fmt.Sprintf("%s_%s_%d_%d_%d.ts", *outPrefix, p.Name, n, k, j)
							muxer = "mpegts"
						} else {
							oname = "-"
							muxer = "null"
						}
						o := ffmpeg.TranscodeOptions{
							Oname:        oname,
							Profile:      p,
							Accel:        accel,
							AudioEncoder: ffmpeg.ComponentOptions{Name: "drop"},
							Muxer:        ffmpeg.ComponentOptions{Name: muxer},
						}
						opts = append(opts, o)
					}
					return opts
				}
				out := profs2opts(profiles)
				t := time.Now()
				_, err := tc.Transcode(in, out)
				end := time.Now()
				fmt.Printf("%s,%d,%d,%0.2v\n", end.Format("2006-01-02 15:04:05.999999999"), k, j, end.Sub(t).Seconds())
				if err != nil {
					panic(err)
				}
				segCount++
			}
			tc.StopTranscoder()
			wg.Done()
		}(i, &wg)
		time.Sleep(300 * time.Millisecond)
	}
	wg.Wait()
	fmt.Fprintf(os.Stderr, "Took %v to transcode %v segments\n",
		time.Now().Sub(start).Seconds(), segCount)
}

func parseVideoProfiles(inp string) []ffmpeg.VideoProfile {
	type profilesJson struct {
		Profiles []struct {
			Name    string `json:"name"`
			Width   int    `json:"width"`
			Height  int    `json:"height"`
			Bitrate int    `json:"bitrate"`
			FPS     uint   `json:"fps"`
			FPSDen  uint   `json:"fpsDen"`
			Profile string `json:"profile"`
			GOP     string `json:"gop"`
		} `json:"profiles"`
	}
	profs := []ffmpeg.VideoProfile{}
	if inp != "" {
		// try opening up json file with profiles
		content, err := ioutil.ReadFile(inp)
		if err == nil && len(content) > 0 {
			// parse json profiles
			resp := &profilesJson{}
			err = json.Unmarshal(content, &resp.Profiles)
			if err != nil {
				panic(err)
			}
			for _, profile := range resp.Profiles {
				name := profile.Name
				if name == "" {
					name = "custom_" + common.DefaultProfileName(
						profile.Width,
						profile.Height,
						profile.Bitrate)
				}
				var gop time.Duration
				if profile.GOP != "" {
					if profile.GOP == "intra" {
						gop = ffmpeg.GOPIntraOnly
					} else {
						gopFloat, err := strconv.ParseFloat(profile.GOP, 64)
						if err != nil {
							panic(err)
						}
						if gopFloat <= 0.0 {
							panic("invalid gop value")
						}
						gop = time.Duration(gopFloat * float64(time.Second))
					}
				}
				encodingProfile, err := common.EncoderProfileNameToValue(profile.Profile)
				if err != nil {
					panic(err)
				}
				prof := ffmpeg.VideoProfile{
					Name:         name,
					Bitrate:      fmt.Sprint(profile.Bitrate),
					Framerate:    profile.FPS,
					FramerateDen: profile.FPSDen,
					Resolution:   fmt.Sprintf("%dx%d", profile.Width, profile.Height),
					Profile:      encodingProfile,
					GOP:          gop,
				}
				profs = append(profs, prof)
			}
		} else {
			// check the built-in profiles
			profs = make([]ffmpeg.VideoProfile, 0)
			presets := strings.Split(inp, ",")
			for _, v := range presets {
				if p, ok := ffmpeg.VideoProfileLookup[strings.TrimSpace(v)]; ok {
					profs = append(profs, p)
				}
			}
		}
		if len(profs) <= 0 {
			panic(fmt.Errorf("No transcoding profiles found"))
		}
	}
	return profs
}
<|endoftext|>


<commit_before>// Copyright 2016-2019 the u-root Authors. All rights reserved
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package main

import (
	"bytes"
	"encoding/hex"
	"fmt"
	"io/ioutil"
	"os"
	"os/exec"
	"path/filepath"
	"strings"
	"testing"

	flag "github.com/spf13/pflag"
)

const (
	testDataDir = "testdata"
)

func resetFlags() {
	*flagFromDump = ""
	*flagType = nil
}

func testOutput(t *testing.T, dumpFile string, args []string, expectedOutFile string) {
	actualOutFile := fmt.Sprintf("%s.actual", expectedOutFile)
	os.Remove(actualOutFile)
	os.Args = []string{os.Args[0], "--from-dump", dumpFile}
	os.Args = append(os.Args, args...)
	flag.Parse()
	defer resetFlags()
	out := &bytes.Buffer{}
	if err := dmiDecode(out); err != nil {
		t.Errorf("%+v %+v %+v: error: %v", dumpFile, args, expectedOutFile, err)
		return
	}
	actualOut := out.Bytes()
	expectedOut, err := ioutil.ReadFile(expectedOutFile)
	if err != nil {
		t.Errorf("%+v %+v %+v: failed to load %s: %v", dumpFile, args, expectedOutFile, expectedOutFile, err)
		return
	}
	if bytes.Compare(actualOut, expectedOut) != 0 {
		ioutil.WriteFile(actualOutFile, actualOut, 0644)
		t.Errorf("%+v %+v %+v: output mismatch, see %s", dumpFile, args, expectedOutFile, actualOutFile)
		diffOut, _ := exec.Command("diff", "-u", expectedOutFile, actualOutFile).CombinedOutput()
		t.Errorf("%+v %+v %+v: diff:\n%s", dumpFile, args, expectedOutFile, string(diffOut))
	}
}

func TestDMIDecode(t *testing.T) {
	bf, err := filepath.Glob("testdata/*.bin")
	if err != nil {
		t.Fatalf("glob failed: %v", err)
	}
	for _, dumpFile := range bf {
		txtFile := strings.TrimSuffix(dumpFile, ".bin") + ".txt"
		testOutput(t, dumpFile, nil, txtFile)
	}
}

func TestDMIDecodeTypeFilters(t *testing.T) {
	testOutput(t, "testdata/Asus-UX307LA.bin", []string{"-t", "system"}, "testdata/Asus-UX307LA.system.txt")
	testOutput(t, "testdata/Asus-UX307LA.bin", []string{"-t", "1,131"}, "testdata/Asus-UX307LA.1_131.txt")
}

func testDumpBin(t *testing.T, entryData, expectedOutData []byte) {
	tmpfile, err := ioutil.TempFile("", "dmidecode")
	if err != nil {
		t.Fatalf("error creating temp file: %v", err)
	}
	tmpfile.Close()
	defer os.Remove(tmpfile.Name())
	textOut := bytes.NewBuffer(nil)
	if err := dumpBin(
		textOut,
		entryData,
		[]byte{0xaa, 0xbb}, // dummy
		tmpfile.Name(),
	); err != nil {
		t.Fatalf("failed to dump bin: %v", err)
	}
	outData, err := ioutil.ReadFile(tmpfile.Name())
	if err != nil {
		t.Fatalf("failed to read output: %v", err)
	}
	if bytes.Compare(outData, expectedOutData) != 0 {
		t.Fatalf("binary data mismatch,\nexpected:\n  %s\ngot:\n  %s", hex.EncodeToString(expectedOutData), hex.EncodeToString(outData))
	}
}

func TestDMIDecodeDumpBin32(t *testing.T) {
	// We expect entry point address to be rewritten and checksum adjusted.
	testDumpBin(
		t,
		[]byte{
			0x5f, 0x53, 0x4d, 0x5f, 0x64, 0x1f, 0x02, 0x08, 0x14, 0x01, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
			0x5f, 0x44, 0x4d, 0x49, 0x5f, 0x37, 0x6e, 0x08, 0x00, 0x50, 0x7c, 0xac, 0x1b, 0x00, 0x28,
		},
		[]byte{
			0x5f, 0x53, 0x4d, 0x5f, 0x64, 0x1f, 0x02, 0x08, 0x14, 0x01, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
			0x5f, 0x44, 0x4d, 0x49, 0x5f, 0x8f, 0x6e, 0x08, 0x20, 0x00, 0x00, 0x00, 0x1b, 0x00, 0x28, 0x00,
			0xaa, 0xbb,
		},
	)
}

func TestDMIDecodeDumpBin64(t *testing.T) {
	// We expect entry point address to be rewritten and checksum adjusted.
	testDumpBin(
		t,
		[]byte{
			0x5f, 0x53, 0x4d, 0x33, 0x5f, 0xe6, 0x18, 0x03, 0x00, 0x00, 0x01, 0x00, 0xe3, 0x0b, 0x00, 0x00,
			0x00, 0xe0, 0x10, 0x8f, 0x00, 0x00, 0x00, 0x00,
		},
		[]byte{
			0x5f, 0x53, 0x4d, 0x33, 0x5f, 0x45, 0x18, 0x03, 0x00, 0x00, 0x01, 0x00, 0xe3, 0x0b, 0x00, 0x00,
			0x20, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
			0xaa, 0xbb,
		},
	)
}
<commit_msg>improve bytes check<commit_after>// Copyright 2016-2019 the u-root Authors. All rights reserved
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package main

import (
	"bytes"
	"encoding/hex"
	"fmt"
	"io/ioutil"
	"os"
	"os/exec"
	"path/filepath"
	"strings"
	"testing"

	flag "github.com/spf13/pflag"
)

const (
	testDataDir = "testdata"
)

func resetFlags() {
	*flagFromDump = ""
	*flagType = nil
}

func testOutput(t *testing.T, dumpFile string, args []string, expectedOutFile string) {
	actualOutFile := fmt.Sprintf("%s.actual", expectedOutFile)
	os.Remove(actualOutFile)
	os.Args = []string{os.Args[0], "--from-dump", dumpFile}
	os.Args = append(os.Args, args...)
	flag.Parse()
	defer resetFlags()
	out := &bytes.Buffer{}
	if err := dmiDecode(out); err != nil {
		t.Errorf("%+v %+v %+v: error: %v", dumpFile, args, expectedOutFile, err)
		return
	}
	actualOut := out.Bytes()
	expectedOut, err := ioutil.ReadFile(expectedOutFile)
	if err != nil {
		t.Errorf("%+v %+v %+v: failed to load %s: %v", dumpFile, args, expectedOutFile, expectedOutFile, err)
		return
	}
	if !bytes.Equal(actualOut, expectedOut) {
		ioutil.WriteFile(actualOutFile, actualOut, 0644)
		t.Errorf("%+v %+v %+v: output mismatch, see %s", dumpFile, args, expectedOutFile, actualOutFile)
		diffOut, _ := exec.Command("diff", "-u", expectedOutFile, actualOutFile).CombinedOutput()
		t.Errorf("%+v %+v %+v: diff:\n%s", dumpFile, args, expectedOutFile, string(diffOut))
	}
}

func TestDMIDecode(t *testing.T) {
	bf, err := filepath.Glob("testdata/*.bin")
	if err != nil {
		t.Fatalf("glob failed: %v", err)
	}
	for _, dumpFile := range bf {
		txtFile := strings.TrimSuffix(dumpFile, ".bin") + ".txt"
		testOutput(t, dumpFile, nil, txtFile)
	}
}

func TestDMIDecodeTypeFilters(t *testing.T) {
	testOutput(t, "testdata/Asus-UX307LA.bin", []string{"-t", "system"}, "testdata/Asus-UX307LA.system.txt")
	testOutput(t, "testdata/Asus-UX307LA.bin", []string{"-t", "1,131"}, "testdata/Asus-UX307LA.1_131.txt")
}

func testDumpBin(t *testing.T, entryData, expectedOutData []byte) {
	tmpfile, err := ioutil.TempFile("", "dmidecode")
	if err != nil {
		t.Fatalf("error creating temp file: %v", err)
	}
	tmpfile.Close()
	defer os.Remove(tmpfile.Name())
	textOut := bytes.NewBuffer(nil)
	if err := dumpBin(
		textOut,
		entryData,
		[]byte{0xaa, 0xbb}, // dummy
		tmpfile.Name(),
	); err != nil {
		t.Fatalf("failed to dump bin: %v", err)
	}
	outData, err := ioutil.ReadFile(tmpfile.Name())
	if err != nil {
		t.Fatalf("failed to read output: %v", err)
	}
	if !bytes.Equal(outData, expectedOutData) {
		t.Fatalf("binary data mismatch,\nexpected:\n  %s\ngot:\n  %s", hex.EncodeToString(expectedOutData), hex.EncodeToString(outData))
	}
}

func TestDMIDecodeDumpBin32(t *testing.T) {
	// We expect entry point address to be rewritten and checksum adjusted.
	testDumpBin(
		t,
		[]byte{
			0x5f, 0x53, 0x4d, 0x5f, 0x64, 0x1f, 0x02, 0x08, 0x14, 0x01, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
			0x5f, 0x44, 0x4d, 0x49, 0x5f, 0x37, 0x6e, 0x08, 0x00, 0x50, 0x7c, 0xac, 0x1b, 0x00, 0x28,
		},
		[]byte{
			0x5f, 0x53, 0x4d, 0x5f, 0x64, 0x1f, 0x02, 0x08, 0x14, 0x01, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
			0x5f, 0x44, 0x4d, 0x49, 0x5f, 0x8f, 0x6e, 0x08, 0x20, 0x00, 0x00, 0x00, 0x1b, 0x00, 0x28, 0x00,
			0xaa, 0xbb,
		},
	)
}

func TestDMIDecodeDumpBin64(t *testing.T) {
	// We expect entry point address to be rewritten and checksum adjusted.
	testDumpBin(
		t,
		[]byte{
			0x5f, 0x53, 0x4d, 0x33, 0x5f, 0xe6, 0x18, 0x03, 0x00, 0x00, 0x01, 0x00, 0xe3, 0x0b, 0x00, 0x00,
			0x00, 0xe0, 0x10, 0x8f, 0x00, 0x00, 0x00, 0x00,
		},
		[]byte{
			0x5f, 0x53, 0x4d, 0x33, 0x5f, 0x45, 0x18, 0x03, 0x00, 0x00, 0x01, 0x00, 0xe3, 0x0b, 0x00, 0x00,
			0x20, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
			0xaa, 0xbb,
		},
	)
}
<|endoftext|>


<commit_before>package dicom

import (
	"fmt"
	"testing"
)

var parser *Parser

func init() {
	parser, _ = NewParser()
}

func TestDefaultDictionary(t *testing.T) {
	_, err := NewParser()
	if err != nil {
		t.Error(err)
	}
}

func TestGetDictEntry(t *testing.T) {
	elem, err := parser.getDictEntry(32736, 16)
	if err != nil {
		t.Error(err)
	}

	if elem.name != "PixelData" {
		t.Errorf("Wrong element name: %s", elem.name)
	}

	if elem.vr != "OX" {
		t.Errorf("Wrong element VR: %s", elem.vr)
	}

}

// TODO: add a test for correctly splitting ranges
func TestSplitTag(t *testing.T) {

	group, element, err := splitTag("(7FE0,0010)")
	if err != nil {
		t.Error(err)
	}

	if group != 0x7FE0 {
		t.Errorf("Error splitting tag. Wrong group: %#x", group)
	}

	if element != 0x0010 {
		t.Errorf("Error splitting tag. Wrong element: %#x", element)
	}

}

func BenchmarkFindMetaGroupLengthTag(b *testing.B) {
	for i := 0; i < b.N; i++ {

		_, err := parser.getDictEntry(2, 0)
		if err != nil {
			fmt.Println(err)
		}

	}
}

func BenchmarkFindPixelDataTag(b *testing.B) {
	for i := 0; i < b.N; i++ {

		_, err := parser.getDictEntry(32736, 16)
		if err != nil {
			fmt.Println(err)
		}

	}
}
<commit_msg>Dictionary tests cleanup<commit_after>package dicom

import (
	"fmt"
	"testing"
)

var parser *Parser

func init() {
	parser, _ = NewParser()
}

func TestDefaultDictionary(t *testing.T) {
	if _, err := NewParser(); err != nil {
		t.Error(err)
	}
}

func TestGetDictEntry(t *testing.T) {
	elem, err := parser.getDictEntry(32736, 16)
	if err != nil {
		t.Error(err)
	}

	if elem.name != "PixelData" {
		t.Errorf("Wrong element name: %s", elem.name)
	}

	if elem.vr != "OX" {
		t.Errorf("Wrong element VR: %s", elem.vr)
	}

}

// TODO: add a test for correctly splitting ranges
func TestSplitTag(t *testing.T) {

	group, element, err := splitTag("(7FE0,0010)")
	if err != nil {
		t.Error(err)
	}

	if group != 0x7FE0 {
		t.Errorf("Error splitting tag. Wrong group: %#x", group)
	}

	if element != 0x0010 {
		t.Errorf("Error splitting tag. Wrong element: %#x", element)
	}

}

func BenchmarkFindMetaGroupLengthTag(b *testing.B) {
	for i := 0; i < b.N; i++ {

		if _, err := parser.getDictEntry(2, 0); err != nil {
			fmt.Println(err)
		}

	}
}

func BenchmarkFindPixelDataTag(b *testing.B) {
	for i := 0; i < b.N; i++ {

		if _, err := parser.getDictEntry(32736, 16); err != nil {
			fmt.Println(err)
		}

	}
}
<|endoftext|>


<commit_before>// Copyright 2015 The Vanadium Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

// Package discovery defines types and interfaces for discovering services.
//
// TODO(jhahn): This is a work in progress and can change without notice.
package discovery

import (
	"v.io/v23/context"
	"v.io/v23/security"
)

// T is the interface for discovery operations; it is the client side library
// for the discovery service.
type T interface {
	Advertiser
	Scanner
	Closer
}

// Advertiser is the interface for advertising services.
type Advertiser interface {
	// Advertise advertises the service to be discovered by "Scanner" implementations.
	// visibility is used to limit the principals that can see the advertisement. An
	// empty set means that there are no restrictions on visibility (i.e, equivalent
	// to []security.BlessingPattern{security.AllPrincipals}). Advertising will continue
	// until the context is canceled or exceeds its deadline and the returned channel
	// will be closed when it stops.
	//
	// It is an error to have simultaneously active advertisements for two identical
	// instances (service.InstanceUuid).
	Advertise(ctx *context.T, service Service, perms []security.BlessingPattern) (<-chan struct{}, error)
}

// AdvertiseCloser is the interface that groups the Advertise and Close methods.
type AdvertiseCloser interface {
	Advertiser
	Closer
}

// Scanner is the interface for scanning services.
type Scanner interface {
	// Scan scans services that match the query and returns the channel on which
	// new discovered services can be read. Scanning will continue until the context
	// is canceled or exceeds its deadline.
	//
	// The query is a WHERE expression of syncQL query against scanned services, where
	// keys are InstanceUuids and values are Service.
	//
	// Examples
	//
	//    v.InstanceName = "v.io/i"
	//    v.InstanceName = "v.io/i" AND v.Attrs["a"] = "v"
	//    v.Attrs["a"] = "v1" OR v.Attrs["a"] = "v2"
	//
	// SyncQL tutorial at:
	//    https://github.com/vanadium/docs/blob/master/tutorials/syncql-tutorial.md
	Scan(ctx *context.T, query string) (<-chan Update, error)
}

// ScanCloser is the interface that groups the Scan and Close methods.
type ScanCloser interface {
	Scanner
	Closer
}

// Closer is the interface that wraps the Close method.
type Closer interface {
	// Close closes all active tasks.
	Close()
}
<commit_msg>discovery: update variable name<commit_after>// Copyright 2015 The Vanadium Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

// Package discovery defines types and interfaces for discovering services.
//
// TODO(jhahn): This is a work in progress and can change without notice.
package discovery

import (
	"v.io/v23/context"
	"v.io/v23/security"
)

// T is the interface for discovery operations; it is the client side library
// for the discovery service.
type T interface {
	Advertiser
	Scanner
	Closer
}

// Advertiser is the interface for advertising services.
type Advertiser interface {
	// Advertise advertises the service to be discovered by "Scanner" implementations.
	// visibility is used to limit the principals that can see the advertisement. An
	// empty set means that there are no restrictions on visibility (i.e, equivalent
	// to []security.BlessingPattern{security.AllPrincipals}). Advertising will continue
	// until the context is canceled or exceeds its deadline and the returned channel
	// will be closed when it stops.
	//
	// It is an error to have simultaneously active advertisements for two identical
	// instances (service.InstanceUuid).
	Advertise(ctx *context.T, service Service, visibility []security.BlessingPattern) (<-chan struct{}, error)
}

// AdvertiseCloser is the interface that groups the Advertise and Close methods.
type AdvertiseCloser interface {
	Advertiser
	Closer
}

// Scanner is the interface for scanning services.
type Scanner interface {
	// Scan scans services that match the query and returns the channel on which
	// new discovered services can be read. Scanning will continue until the context
	// is canceled or exceeds its deadline.
	//
	// The query is a WHERE expression of syncQL query against scanned services, where
	// keys are InstanceUuids and values are Service.
	//
	// Examples
	//
	//    v.InstanceName = "v.io/i"
	//    v.InstanceName = "v.io/i" AND v.Attrs["a"] = "v"
	//    v.Attrs["a"] = "v1" OR v.Attrs["a"] = "v2"
	//
	// SyncQL tutorial at:
	//    https://github.com/vanadium/docs/blob/master/tutorials/syncql-tutorial.md
	Scan(ctx *context.T, query string) (<-chan Update, error)
}

// ScanCloser is the interface that groups the Scan and Close methods.
type ScanCloser interface {
	Scanner
	Closer
}

// Closer is the interface that wraps the Close method.
type Closer interface {
	// Close closes all active tasks.
	Close()
}
<|endoftext|>


<commit_before>package keycreator

import (
	"errors"
	"strings"
	"time"

	"github.com/dgrijalva/jwt-go"
)

var (
	ErrSignUsernameEmpty   = errors.New("Username is empty")
	ErrSignKontrolURLEmpty = errors.New("Kontrol URL is empty")
	ErrSignPrivateKeyEmpty = errors.New("Private key is empty")
	ErrSignPublicKeyEmpty  = errors.New("Public key is empty")
)

type Key struct {
	KontrolURL        string
	KontrolPrivateKey string
	KontrolPublicKey  string
}

// Create signs a new key and returns the token back
func (k *Key) Create(username, kiteId string) (string, error) {
	if username == "" {
		return "", ErrSignUsernameEmpty
	}

	if k.KontrolURL == "" {
		return "", ErrSignKontrolURLEmpty
	}

	if k.KontrolPrivateKey == "" {
		return "", ErrSignPrivateKeyEmpty
	}

	if k.KontrolPublicKey == "" {
		return "", ErrSignPublicKeyEmpty
	}

	token := jwt.New(jwt.GetSigningMethod("RS256"))

	token.Claims = map[string]interface{}{
		"iss":        "koding",                              // Issuer, should be the same username as kontrol
		"sub":        username,                              // Subject
		"iat":        time.Now().UTC().Unix(),               // Issued At
		"jti":        kiteId,                                // JWT ID
		"kontrolURL": k.KontrolURL,                          // Kontrol URL
		"kontrolKey": strings.TrimSpace(k.KontrolPublicKey), // Public key of kontrol
	}

	return token.SignedString([]byte(k.KontrolPrivateKey))
}
<commit_msg>keycreator: update to jwt-go 3.0 api<commit_after>package keycreator

import (
	"errors"
	"strings"
	"time"

	"github.com/dgrijalva/jwt-go"
	"github.com/koding/kite/kitekey"
)

var (
	ErrSignUsernameEmpty   = errors.New("Username is empty")
	ErrSignKontrolURLEmpty = errors.New("Kontrol URL is empty")
	ErrSignPrivateKeyEmpty = errors.New("Private key is empty")
	ErrSignPublicKeyEmpty  = errors.New("Public key is empty")
)

type Key struct {
	KontrolURL        string
	KontrolPrivateKey string
	KontrolPublicKey  string
}

// Create signs a new key and returns the token back
func (k *Key) Create(username, kiteId string) (string, error) {
	if username == "" {
		return "", ErrSignUsernameEmpty
	}

	if k.KontrolURL == "" {
		return "", ErrSignKontrolURLEmpty
	}

	if k.KontrolPrivateKey == "" {
		return "", ErrSignPrivateKeyEmpty
	}

	rsaKey, err := jwt.ParseRSAPrivateKeyFromPEM([]byte(k.KontrolPrivateKey))
	if err != nil {
		return "", err
	}

	if k.KontrolPublicKey == "" {
		return "", ErrSignPublicKeyEmpty
	}

	claims := &kitekey.KiteClaims{
		StandardClaims: jwt.StandardClaims{
			Issuer:   "koding",
			Subject:  username,
			IssuedAt: time.Now().UTC().Unix(),
			Id:       kiteId,
		},
		KontrolURL: k.KontrolURL,
		KontrolKey: strings.TrimSpace(k.KontrolPublicKey),
	}

	token := jwt.NewWithClaims(jwt.GetSigningMethod("RS256"), claims)

	return token.SignedString(rsaKey)
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"github.com/gorilla/sessions"
	"github.com/nranchev/go-libGeoIP"
	"html/template"
	"io"
	"koding/kontrol/kontrolhelper"
	"koding/kontrol/kontrolproxy/proxyconfig"
	"koding/tools/config"
	"log"
	"net"
	"net/http"
	"net/url"
	"strconv"
	"strings"
	"time"
)

func init() {
	log.SetPrefix("kontrol-proxy ")
}

type RabbitChannel struct {
	ReplyTo string
	Receive chan []byte
}

var proxyDB *proxyconfig.ProxyConfiguration
var amqpStream *AmqpStream
var connections map[string]RabbitChannel
var geoIP *libgeo.GeoIP
var hostname = kontrolhelper.CustomHostname()
var store = sessions.NewCookieStore([]byte("kontrolproxy-secret-key"))
var templates = template.Must(template.ParseFiles("go/templates//proxy/securepage.html"))

func main() {
	log.Printf("kontrol proxy started ")
	connections = make(map[string]RabbitChannel)

	// open and read from DB
	var err error
	proxyDB, err = proxyconfig.Connect()
	if err != nil {
		log.Fatalf("proxyconfig mongodb connect: %s", err)
	}

	err = proxyDB.AddProxy(hostname)
	if err != nil {
		log.Println(err)
	}

	// load GeoIP db into memory
	dbFile := "GeoIP.dat"
	geoIP, err = libgeo.Load(dbFile)
	if err != nil {
		log.Printf("load GeoIP.dat: %s\n", err.Error())
	}

	// create amqpStream for rabbitmq proxyieng
	amqpStream = setupAmqp()

	reverseProxy := &ReverseProxy{}
	// http.HandleFunc("/", reverseProxy.ServeHTTP) this works for 1.1
	http.HandleFunc("/", func(w http.ResponseWriter, r *http.Request) {
		reverseProxy.ServeHTTP(w, r)
	})

	port := strconv.Itoa(config.Current.Kontrold.Proxy.Port)
	portssl := strconv.Itoa(config.Current.Kontrold.Proxy.PortSSL)
	sslips := strings.Split(config.Current.Kontrold.Proxy.SSLIPS, ",")

	for _, sslip := range sslips {
		go func(sslip string) {
			err = http.ListenAndServeTLS(sslip+":"+portssl, sslip+"_cert.pem", sslip+"_key.pem", nil)
			if err != nil {
				log.Printf("https mode is disabled. please add cert.pem and key.pem files. %s %s", err, sslip)
			} else {
				log.Printf("https mode is enabled. serving at :%s ...", portssl)
			}
		}(sslip)
	}

	log.Printf("normal mode is enabled. serving at :%s ...", port)
	err = http.ListenAndServe(":"+port, nil)
	if err != nil {
		log.Println(err)
	}
}

/*************************************************
*
*  util functions
*
*  - arslan
*************************************************/
// Given a string of the form "host", "host:port", or "[ipv6::address]:port",
// return true if the string includes a port.
func hasPort(s string) bool { return strings.LastIndex(s, ":") > strings.LastIndex(s, "]") }

// Given a string of the form "host", "port", returns "host:port"
func addPort(host, port string) string {
	if ok := hasPort(host); ok {
		return host
	}

	return host + ":" + port
}

// Check if a server is alive or not
func checkServer(host string) error {
	c, err := net.Dial("tcp", host)
	if err != nil {
		return err
	}
	c.Close()
	return nil
}

/*************************************************
*
*  modified version of go's reverseProxy source code
*  has support for dynamic target url, websockets and amqp
*
*  - arslan
*************************************************/

// ReverseProxy is an HTTP Handler that takes an incoming request and
// sends it to another server, proxying the response back to the
// client.
type ReverseProxy struct {
	// The transport used to perform proxy requests.
	// If nil, http.DefaultTransport is used.
	Transport http.RoundTripper
}

func singleJoiningSlash(a, b string) string {
	aslash := strings.HasSuffix(a, "/")
	bslash := strings.HasPrefix(b, "/")
	switch {
	case aslash && bslash:
		return a + b[1:]
	case !aslash && !bslash:
		return a + "/" + b
	}
	return a + b
}

func copyHeader(dst, src http.Header) {
	for k, vv := range src {
		for _, v := range vv {
			dst.Add(k, v)
		}
	}
}

// Hop-by-hop headers. These are removed when sent to the backend.
// http://www.w3.org/Protocols/rfc2616/rfc2616-sec13.html
var hopHeaders = []string{
	"Connection",
	"Keep-Alive",
	"Proxy-Authenticate",
	"Proxy-Authorization",
	"Te", // canonicalized version of "TE"
	"Trailers",
	"Transfer-Encoding",
	"Upgrade",
}

func (p *ReverseProxy) ServeHTTP(rw http.ResponseWriter, req *http.Request) {
	// redirect http to https
	if req.TLS == nil && req.Host == "new.koding.com" {
		http.Redirect(rw, req, "https://new.koding.com"+req.RequestURI, http.StatusMovedPermanently)
		return
	}

	// Display error when someone hits the main page
	if hostname == req.Host {
		io.WriteString(rw, "Hello kontrol proxy :)")
		return
	}

	outreq := new(http.Request)
	*outreq = *req // includes shallow copies of maps, but okay

	user, err := populateUser(outreq)
	if err != nil {
		log.Printf("\nWARNING: parsing incoming request %s: %s", outreq.Host, err.Error())
		io.WriteString(rw, fmt.Sprintf("{\"err\":\"%s\"}\n", err.Error()))
		return
	}

	target := user.Target
	if user.Domain.LoadBalancer.Mode == "sticky" {
		sessionName := fmt.Sprintf("kodingproxy-%s-%s", outreq.Host, user.IP)
		session, _ := store.Get(req, sessionName)
		targetURL, ok := session.Values["GOSESSIONID"]
		if ok {
			fmt.Printf("proxy via session cookie\t: %s --> %s\n", user.Domain.Domain, user.Target.Host)
			target, err = url.Parse(targetURL.(string))
			if err != nil {
				io.WriteString(rw, fmt.Sprintf("{\"err\":\"%s\"}\n", err.Error()))
				return
			}
		} else {
			fmt.Printf("proxy via db\t: %s --> %s\n", user.Domain.Domain, user.Target.Host)
			session.Values["GOSESSIONID"] = target.String()
			session.Save(outreq, rw)
		}
	}

	if user.Redirect {
		http.Redirect(rw, req, user.Target.String(), http.StatusTemporaryRedirect)
		return
	}

	_, err = validate(user)
	if err != nil {
		if err == ErrSecurePage {
			sessionName := fmt.Sprintf("kodingproxy-%s-%s", outreq.Host, user.IP)
			// We're ignoring the error resulted from decoding an existing
			// session: Get() always returns a session, even if empty.
			session, _ := store.Get(req, sessionName)

			// Timeout for secure page. After timeout secure page is showed
			// again to the user
			session.Options = &sessions.Options{MaxAge: 20} //seconds

			_, ok := session.Values["securePage"]
			if !ok {
				session.Values["securePage"] = time.Now().String()
				session.Save(req, rw)
				err := templates.ExecuteTemplate(rw, "securepage.html", user)
				if err != nil {
					http.Error(rw, err.Error(), http.StatusInternalServerError)
				}
				return
			}
		} else {
			log.Printf("error validating user %s: %s", user.IP, err.Error())
			io.WriteString(rw, fmt.Sprintf("{\"err\":\"%s\"}\n", err.Error()))
			return

		}
	}

	// Smart handling incoming request path/query, example:
	// incoming : foo.com/dir
	// target	: bar.com/base
	// proxy to : bar.com/base/dir
	outreq.URL.Scheme = target.Scheme
	outreq.URL.Host = target.Host
	outreq.URL.Path = singleJoiningSlash(target.Path, outreq.URL.Path)

	// incoming : foo.com/name=arslan
	// target	: bar.com/q=example
	// proxy to : bar.com/q=example&name=arslan
	if target.RawQuery == "" || outreq.URL.RawQuery == "" {
		outreq.URL.RawQuery = target.RawQuery + outreq.URL.RawQuery
	} else {
		outreq.URL.RawQuery = target.RawQuery + "&" + outreq.URL.RawQuery
	}

	outreq.Proto = "HTTP/1.1"
	outreq.ProtoMajor = 1
	outreq.ProtoMinor = 1
	outreq.Close = false

	// if connection is of type websocket, hijacking is used instead of http proxy
	// https://groups.google.com/d/msg/golang-nuts/KBx9pDlvFOc/edt4iad96nwJ
	if isWebsocket(outreq) {
		rConn, err := net.Dial("tcp", outreq.URL.Host)
		if err != nil {
			http.Error(rw, "Error contacting backend server.", http.StatusInternalServerError)
			log.Printf("Error dialing websocket backend %s: %v", outreq.URL.Host, err)
			return
		}

		hj, ok := rw.(http.Hijacker)
		if !ok {
			http.Error(rw, "Not a hijacker?", http.StatusInternalServerError)
			return
		}

		conn, _, err := hj.Hijack()
		if err != nil {
			log.Printf("Hijack error: %v", err)
			return
		}
		defer conn.Close()
		defer rConn.Close()

		err = req.Write(rConn)
		if err != nil {
			log.Printf("Error copying request to target: %v", err)
			return
		}

		go p.copyResponse(rConn, conn)
		p.copyResponse(conn, rConn)

	} else {
		go logDomainRequests(outreq.Host)
		go logProxyStat(hostname, user.Country)

		transport := p.Transport
		if transport == nil {
			transport = http.DefaultTransport
		}

		// Remove hop-by-hop headers to the backend.  Especially
		// important is "Connection" because we want a persistent
		// connection, regardless of what the client sent to us.  This
		// is modifying the same underlying map from req (shallow
		// copied above) so we only copy it if necessary.
		copiedHeaders := false
		for _, h := range hopHeaders {
			if outreq.Header.Get(h) != "" {
				if !copiedHeaders {
					outreq.Header = make(http.Header)
					copyHeader(outreq.Header, req.Header)
					copiedHeaders = true
				}
				outreq.Header.Del(h)
			}
		}

		if clientIP, _, err := net.SplitHostPort(req.RemoteAddr); err == nil {
			// If we aren't the first proxy retain prior
			// X-Forwarded-For information as a comma+space
			// separated list and fold multiple headers into one.
			if prior, ok := outreq.Header["X-Forwarded-For"]; ok {
				clientIP = strings.Join(prior, ", ") + ", " + clientIP
			}
			outreq.Header.Set("X-Forwarded-For", clientIP)
		}

		res := new(http.Response)

		if !hasPort(outreq.URL.Host) {
			outreq.URL.Host = addPort(outreq.URL.Host, "80")
		}

		res, err = transport.RoundTrip(outreq)
		if err != nil {
			io.WriteString(rw, fmt.Sprint(err))
			return
		}
		defer res.Body.Close()

		// rabbitKey, err := lookupRabbitKey(user.Username, user.Servicename, user.Key)
		// if err != nil {
		// 	// add :80 if not available
		// } else {
		// 	fmt.Println("connection via rabbitmq")
		// 	res, err = rabbitTransport(outreq, user, rabbitKey)
		// 	if err != nil {
		// 		log.Printf("rabbit proxy %s", err.Error())
		// 		io.WriteString(rw, fmt.Sprintf("{\"err\":\"%s\"}\n", err.Error()))
		// 		return
		// 	}

		// }

		copyHeader(rw.Header(), res.Header)
		rw.WriteHeader(res.StatusCode)
		p.copyResponse(rw, res.Body)
		return
	}

}

func (p *ReverseProxy) copyResponse(dst io.Writer, src io.Reader) {
	io.Copy(dst, src)
}
<commit_msg>kontrolproxy: store request for unique IP's every one hour<commit_after>package main

import (
	"fmt"
	"github.com/gorilla/sessions"
	"github.com/nranchev/go-libGeoIP"
	"html/template"
	"io"
	"koding/kontrol/kontrolhelper"
	"koding/kontrol/kontrolproxy/proxyconfig"
	"koding/tools/config"
	"log"
	"net"
	"net/http"
	"net/url"
	"strconv"
	"strings"
	"sync"
	"time"
)

func init() {
	log.SetPrefix("kontrol-proxy ")
}

type RabbitChannel struct {
	ReplyTo string
	Receive chan []byte
}

var proxyDB *proxyconfig.ProxyConfiguration
var amqpStream *AmqpStream
var connections = make(map[string]RabbitChannel)
var geoIP *libgeo.GeoIP
var hostname = kontrolhelper.CustomHostname()
var store = sessions.NewCookieStore([]byte("kontrolproxy-secret-key"))
var templates = template.Must(template.ParseFiles("go/templates//proxy/securepage.html"))
var users = make(map[string]time.Time)
var usersLock sync.RWMutex

func main() {
	log.Printf("kontrol proxy started ")
	// open and read from DB
	var err error
	proxyDB, err = proxyconfig.Connect()
	if err != nil {
		log.Fatalf("proxyconfig mongodb connect: %s", err)
	}

	err = proxyDB.AddProxy(hostname)
	if err != nil {
		log.Println(err)
	}

	// load GeoIP db into memory
	dbFile := "GeoIP.dat"
	geoIP, err = libgeo.Load(dbFile)
	if err != nil {
		log.Printf("load GeoIP.dat: %s\n", err.Error())
	}

	// create amqpStream for rabbitmq proxyieng
	amqpStream = setupAmqp()

	reverseProxy := &ReverseProxy{}
	// http.HandleFunc("/", reverseProxy.ServeHTTP) this works for 1.1
	http.HandleFunc("/", func(w http.ResponseWriter, r *http.Request) {
		reverseProxy.ServeHTTP(w, r)
	})

	port := strconv.Itoa(config.Current.Kontrold.Proxy.Port)
	portssl := strconv.Itoa(config.Current.Kontrold.Proxy.PortSSL)
	sslips := strings.Split(config.Current.Kontrold.Proxy.SSLIPS, ",")

	for _, sslip := range sslips {
		go func(sslip string) {
			err = http.ListenAndServeTLS(sslip+":"+portssl, sslip+"_cert.pem", sslip+"_key.pem", nil)
			if err != nil {
				log.Printf("https mode is disabled. please add cert.pem and key.pem files. %s %s", err, sslip)
			} else {
				log.Printf("https mode is enabled. serving at :%s ...", portssl)
			}
		}(sslip)
	}

	log.Printf("normal mode is enabled. serving at :%s ...", port)
	err = http.ListenAndServe(":"+port, nil)
	if err != nil {
		log.Println(err)
	}
}

/*************************************************
*
*  util functions
*
*  - arslan
*************************************************/
// Given a string of the form "host", "host:port", or "[ipv6::address]:port",
// return true if the string includes a port.
func hasPort(s string) bool { return strings.LastIndex(s, ":") > strings.LastIndex(s, "]") }

// Given a string of the form "host", "port", returns "host:port"
func addPort(host, port string) string {
	if ok := hasPort(host); ok {
		return host
	}

	return host + ":" + port
}

// Check if a server is alive or not
func checkServer(host string) error {
	c, err := net.Dial("tcp", host)
	if err != nil {
		return err
	}
	c.Close()
	return nil
}

/*************************************************
*
*  modified version of go's reverseProxy source code
*  has support for dynamic target url, websockets and amqp
*
*  - arslan
*************************************************/

// ReverseProxy is an HTTP Handler that takes an incoming request and
// sends it to another server, proxying the response back to the
// client.
type ReverseProxy struct {
	// The transport used to perform proxy requests.
	// If nil, http.DefaultTransport is used.
	Transport http.RoundTripper
}

func singleJoiningSlash(a, b string) string {
	aslash := strings.HasSuffix(a, "/")
	bslash := strings.HasPrefix(b, "/")
	switch {
	case aslash && bslash:
		return a + b[1:]
	case !aslash && !bslash:
		return a + "/" + b
	}
	return a + b
}

func copyHeader(dst, src http.Header) {
	for k, vv := range src {
		for _, v := range vv {
			dst.Add(k, v)
		}
	}
}

// Hop-by-hop headers. These are removed when sent to the backend.
// http://www.w3.org/Protocols/rfc2616/rfc2616-sec13.html
var hopHeaders = []string{
	"Connection",
	"Keep-Alive",
	"Proxy-Authenticate",
	"Proxy-Authorization",
	"Te", // canonicalized version of "TE"
	"Trailers",
	"Transfer-Encoding",
	"Upgrade",
}

func (p *ReverseProxy) ServeHTTP(rw http.ResponseWriter, req *http.Request) {
	// redirect http to https
	if req.TLS == nil && req.Host == "new.koding.com" {
		http.Redirect(rw, req, "https://new.koding.com"+req.RequestURI, http.StatusMovedPermanently)
		return
	}

	// Display error when someone hits the main page
	if hostname == req.Host {
		io.WriteString(rw, "Hello kontrol proxy :)")
		return
	}

	outreq := new(http.Request)
	*outreq = *req // includes shallow copies of maps, but okay

	user, err := populateUser(outreq)
	if err != nil {
		log.Printf("\nWARNING: parsing incoming request %s: %s", outreq.Host, err.Error())
		io.WriteString(rw, fmt.Sprintf("{\"err\":\"%s\"}\n", err.Error()))
		return
	}

	target := user.Target
	if user.Domain.LoadBalancer.Mode == "sticky" {
		sessionName := fmt.Sprintf("kodingproxy-%s-%s", outreq.Host, user.IP)
		session, _ := store.Get(req, sessionName)
		targetURL, ok := session.Values["GOSESSIONID"]
		if ok {
			fmt.Printf("proxy via session cookie\t: %s --> %s\n", user.Domain.Domain, user.Target.Host)
			target, err = url.Parse(targetURL.(string))
			if err != nil {
				io.WriteString(rw, fmt.Sprintf("{\"err\":\"%s\"}\n", err.Error()))
				return
			}
		} else {
			fmt.Printf("proxy via db\t: %s --> %s\n", user.Domain.Domain, user.Target.Host)
			session.Values["GOSESSIONID"] = target.String()
			session.Save(outreq, rw)
		}
	}

	if user.Redirect {
		http.Redirect(rw, req, user.Target.String(), http.StatusTemporaryRedirect)
		return
	}

	_, err = validate(user)
	if err != nil {
		if err == ErrSecurePage {
			sessionName := fmt.Sprintf("kodingproxy-%s-%s", outreq.Host, user.IP)
			// We're ignoring the error resulted from decoding an existing
			// session: Get() always returns a session, even if empty.
			session, _ := store.Get(req, sessionName)

			// Timeout for secure page. After timeout secure page is showed
			// again to the user
			session.Options = &sessions.Options{MaxAge: 20} //seconds

			_, ok := session.Values["securePage"]
			if !ok {
				session.Values["securePage"] = time.Now().String()
				session.Save(req, rw)
				err := templates.ExecuteTemplate(rw, "securepage.html", user)
				if err != nil {
					http.Error(rw, err.Error(), http.StatusInternalServerError)
				}
				return
			}
		} else {
			log.Printf("error validating user %s: %s", user.IP, err.Error())
			io.WriteString(rw, fmt.Sprintf("{\"err\":\"%s\"}\n", err.Error()))
			return

		}
	}

	// Smart handling incoming request path/query, example:
	// incoming : foo.com/dir
	// target	: bar.com/base
	// proxy to : bar.com/base/dir
	outreq.URL.Scheme = target.Scheme
	outreq.URL.Host = target.Host
	outreq.URL.Path = singleJoiningSlash(target.Path, outreq.URL.Path)

	// incoming : foo.com/name=arslan
	// target	: bar.com/q=example
	// proxy to : bar.com/q=example&name=arslan
	if target.RawQuery == "" || outreq.URL.RawQuery == "" {
		outreq.URL.RawQuery = target.RawQuery + outreq.URL.RawQuery
	} else {
		outreq.URL.RawQuery = target.RawQuery + "&" + outreq.URL.RawQuery
	}

	outreq.Proto = "HTTP/1.1"
	outreq.ProtoMajor = 1
	outreq.ProtoMinor = 1
	outreq.Close = false

	if !isUserRegistered(user.IP) {
		go registerUser(user.IP)
		go logDomainRequests(outreq.Host)
		go logProxyStat(hostname, user.Country)
	}

	// if connection is of type websocket, hijacking is used instead of http proxy
	// https://groups.google.com/d/msg/golang-nuts/KBx9pDlvFOc/edt4iad96nwJ
	if isWebsocket(outreq) {
		rConn, err := net.Dial("tcp", outreq.URL.Host)
		if err != nil {
			http.Error(rw, "Error contacting backend server.", http.StatusInternalServerError)
			log.Printf("Error dialing websocket backend %s: %v", outreq.URL.Host, err)
			return
		}

		hj, ok := rw.(http.Hijacker)
		if !ok {
			http.Error(rw, "Not a hijacker?", http.StatusInternalServerError)
			return
		}

		conn, _, err := hj.Hijack()
		if err != nil {
			log.Printf("Hijack error: %v", err)
			return
		}
		defer conn.Close()
		defer rConn.Close()

		err = req.Write(rConn)
		if err != nil {
			log.Printf("Error copying request to target: %v", err)
			return
		}

		go p.copyResponse(rConn, conn)
		p.copyResponse(conn, rConn)

	} else {

		transport := p.Transport
		if transport == nil {
			transport = http.DefaultTransport
		}

		// Remove hop-by-hop headers to the backend.  Especially
		// important is "Connection" because we want a persistent
		// connection, regardless of what the client sent to us.  This
		// is modifying the same underlying map from req (shallow
		// copied above) so we only copy it if necessary.
		copiedHeaders := false
		for _, h := range hopHeaders {
			if outreq.Header.Get(h) != "" {
				if !copiedHeaders {
					outreq.Header = make(http.Header)
					copyHeader(outreq.Header, req.Header)
					copiedHeaders = true
				}
				outreq.Header.Del(h)
			}
		}

		if clientIP, _, err := net.SplitHostPort(req.RemoteAddr); err == nil {
			// If we aren't the first proxy retain prior
			// X-Forwarded-For information as a comma+space
			// separated list and fold multiple headers into one.
			if prior, ok := outreq.Header["X-Forwarded-For"]; ok {
				clientIP = strings.Join(prior, ", ") + ", " + clientIP
			}
			outreq.Header.Set("X-Forwarded-For", clientIP)
		}

		res := new(http.Response)

		if !hasPort(outreq.URL.Host) {
			outreq.URL.Host = addPort(outreq.URL.Host, "80")
		}

		res, err = transport.RoundTrip(outreq)
		if err != nil {
			io.WriteString(rw, fmt.Sprint(err))
			return
		}
		defer res.Body.Close()

		// rabbitKey, err := lookupRabbitKey(user.Username, user.Servicename, user.Key)
		// if err != nil {
		// 	// add :80 if not available
		// } else {
		// 	fmt.Println("connection via rabbitmq")
		// 	res, err = rabbitTransport(outreq, user, rabbitKey)
		// 	if err != nil {
		// 		log.Printf("rabbit proxy %s", err.Error())
		// 		io.WriteString(rw, fmt.Sprintf("{\"err\":\"%s\"}\n", err.Error()))
		// 		return
		// 	}

		// }

		copyHeader(rw.Header(), res.Header)
		rw.WriteHeader(res.StatusCode)
		p.copyResponse(rw, res.Body)
		return
	}

}

func (p *ReverseProxy) copyResponse(dst io.Writer, src io.Reader) {
	io.Copy(dst, src)
}

func registerUser(ip string) {
	usersLock.Lock()
	defer usersLock.Unlock()
	users[ip] = time.Now()
	if len(users) == 1 {
		go cleaner()
	}
}

// The goroutine basically does this: as long as there are users in the map, it
// finds the one it should be deleted next, sleeps until it's time to delete it
// (one hour - time since user registration) and deletes it.  If there are no
// users, the goroutine exits and a new one is created the next time a user is
// registered. The time.Sleep goes toward zero, thus it will not lock the
// for iterator forever.
func cleaner() {
	usersLock.RLock()
	for len(users) > 0 {
		var nextTime time.Time
		var nextUser string
		for u, t := range users {
			if nextTime.IsZero() || t.Before(nextTime) {
				nextTime = t
				nextUser = u
			}
		}
		usersLock.RUnlock()
		// negative duration is no-op, means it will not panic
		time.Sleep(time.Hour - time.Now().Sub(nextTime))
		usersLock.Lock()
		delete(users, nextUser)
		usersLock.Unlock()
		usersLock.RLock()
	}
	usersLock.RUnlock()
}

// Needed to avoid race condition between multiple go routines
func isUserRegistered(ip string) bool {
	usersLock.RLock()
	defer usersLock.RUnlock()
	_, ok := users[ip]
	return ok
}
<|endoftext|>


<commit_before>// Copyright 2012, 2013 Canonical Ltd.
// Licensed under the AGPLv3, see LICENCE file for details.

package state

import (
	"crypto/tls"
	"crypto/x509"
	stderrors "errors"
	"fmt"
	"net"
	"strings"
	"time"

	"labix.org/v2/mgo"
	"labix.org/v2/mgo/txn"

	"launchpad.net/juju-core/cert"
	"launchpad.net/juju-core/constraints"
	"launchpad.net/juju-core/environs/config"
	"launchpad.net/juju-core/errors"
	"launchpad.net/juju-core/state/presence"
	"launchpad.net/juju-core/state/watcher"
	"launchpad.net/juju-core/utils"
)

// Info encapsulates information about cluster of
// servers holding juju state and can be used to make a
// connection to that cluster.
type Info struct {
	// Addrs gives the addresses of the MongoDB servers for the state.
	// Each address should be in the form address:port.
	Addrs []string

	// CACert holds the CA certificate that will be used
	// to validate the state server's certificate, in PEM format.
	CACert []byte

	// Tag holds the name of the entity that is connecting.
	// It should be empty when connecting as an administrator.
	Tag string

	// Password holds the password for the connecting entity.
	Password string
}

// DialOpts holds configuration parameters that control the
// Dialing behavior when connecting to a state server.
type DialOpts struct {
	// Timeout is the amount of time to wait contacting
	// a state server.
	Timeout time.Duration
}

// DefaultDialOpts returns a DialOpts representing the default
// parameters for contacting a state server.
func DefaultDialOpts() DialOpts {
	return DialOpts{
		Timeout: 30 * time.Second,
	}
}

// Open connects to the server described by the given
// info, waits for it to be initialized, and returns a new State
// representing the environment connected to.
// It returns unauthorizedError if access is unauthorized.
func Open(info *Info, opts DialOpts) (*State, error) {
	logger.Infof("opening state; mongo addresses: %q; entity %q", info.Addrs, info.Tag)
	if len(info.Addrs) == 0 {
		return nil, stderrors.New("no mongo addresses")
	}
	if len(info.CACert) == 0 {
		return nil, stderrors.New("missing CA certificate")
	}
	xcert, err := cert.ParseCert(info.CACert)
	if err != nil {
		return nil, fmt.Errorf("cannot parse CA certificate: %v", err)
	}
	pool := x509.NewCertPool()
	pool.AddCert(xcert)
	tlsConfig := &tls.Config{
		RootCAs:    pool,
		ServerName: "anything",
	}
	dial := func(addr net.Addr) (net.Conn, error) {
		c, err := net.Dial("tcp", addr.String())
		if err != nil {
			logger.Debugf("connection failed, will retry: %v", err)
			return nil, err
		}
		cc := tls.Client(c, tlsConfig)
		if err := cc.Handshake(); err != nil {
			logger.Errorf("TLS handshake failed: %v", err)
			return nil, err
		}
		return cc, nil
	}
	session, err := mgo.DialWithInfo(&mgo.DialInfo{
		Addrs:   info.Addrs,
		Timeout: opts.Timeout,
		Dial:    dial,
	})
	if err != nil {
		return nil, err
	}
	logger.Infof("connection established")
	st, err := newState(session, info)
	if err != nil {
		session.Close()
		return nil, err
	}
	return st, nil
}

// Initialize sets up an initial empty state and returns it.
// This needs to be performed only once for a given environment.
// It returns unauthorizedError if access is unauthorized.
func Initialize(info *Info, cfg *config.Config, opts DialOpts) (rst *State, err error) {
	st, err := Open(info, opts)
	if err != nil {
		return nil, err
	}
	defer func() {
		if err != nil {
			st.Close()
		}
	}()
	// A valid environment is used as a signal that the
	// state has already been initalized. If this is the case
	// do nothing.
	if _, err := st.Environment(); err == nil {
		return st, nil
	} else if !errors.IsNotFoundError(err) {
		return nil, err
	}
	logger.Infof("initializing environment")
	if err := checkEnvironConfig(cfg); err != nil {
		return nil, err
	}
	uuid, err := utils.NewUUID()
	if err != nil {
		return nil, fmt.Errorf("environment UUID cannot be created: %v", err)
	}
	ops := []txn.Op{
		createConstraintsOp(st, environGlobalKey, constraints.Value{}),
		createSettingsOp(st, environGlobalKey, cfg.AllAttrs()),
		createEnvironmentOp(st, cfg.Name(), uuid.String()),
	}
	if err := st.runTransaction(ops); err == txn.ErrAborted {
		// The config was created in the meantime.
		return st, nil
	} else if err != nil {
		return nil, err
	}
	return st, nil
}

var indexes = []struct {
	collection string
	key        []string
}{
	// After the first public release, do not remove entries from here
	// without adding them to a list of indexes to drop, to ensure
	// old databases are modified to have the correct indexes.
	{"relations", []string{"endpoints.relationname"}},
	{"relations", []string{"endpoints.servicename"}},
	{"units", []string{"service"}},
	{"units", []string{"principal"}},
	{"units", []string{"machineid"}},
	{"users", []string{"name"}},
}

// The capped collection used for transaction logs defaults to 10MB.
// It's tweaked in export_test.go to 1MB to avoid the overhead of
// creating and deleting the large file repeatedly in tests.
var (
	logSize      = 10000000
	logSizeTests = 1000000
)

func maybeUnauthorized(err error, msg string) error {
	if err == nil {
		return nil
	}
	// Unauthorized access errors have no error code,
	// just a simple error string.
	if err.Error() == "auth fails" {
		return errors.NewUnauthorizedError(err, msg)
	}
	if err, ok := err.(*mgo.QueryError); ok && err.Code == 10057 {
		return errors.NewUnauthorizedError(err, msg)
	}
	return fmt.Errorf("%s: %v", msg, err)
}

func newState(session *mgo.Session, info *Info) (*State, error) {
	db := session.DB("juju")
	pdb := session.DB("presence")
	if info.Tag != "" {
		if err := db.Login(info.Tag, info.Password); err != nil {
			return nil, maybeUnauthorized(err, fmt.Sprintf("cannot log in to juju database as %q", info.Tag))
		}
		if err := pdb.Login(info.Tag, info.Password); err != nil {
			return nil, maybeUnauthorized(err, fmt.Sprintf("cannot log in to presence database as %q", info.Tag))
		}
	} else if info.Password != "" {
		admin := session.DB("admin")
		if err := admin.Login("admin", info.Password); err != nil {
			return nil, maybeUnauthorized(err, "cannot log in to admin database")
		}
	}
	st := &State{
		info:           info,
		db:             db,
		environments:   db.C("environments"),
		charms:         db.C("charms"),
		machines:       db.C("machines"),
		containerRefs:  db.C("containerRefs"),
		instanceData:   db.C("instanceData"),
		relations:      db.C("relations"),
		relationScopes: db.C("relationscopes"),
		services:       db.C("services"),
		minUnits:       db.C("minunits"),
		settings:       db.C("settings"),
		settingsrefs:   db.C("settingsrefs"),
		constraints:    db.C("constraints"),
		units:          db.C("units"),
		users:          db.C("users"),
		presence:       pdb.C("presence"),
		cleanups:       db.C("cleanups"),
		annotations:    db.C("annotations"),
		statuses:       db.C("statuses"),
	}
	log := db.C("txns.log")
	logInfo := mgo.CollectionInfo{Capped: true, MaxBytes: logSize}
	// The lack of error code for this error was reported upstream:
	//     https://jira.klmongodb.org/browse/SERVER-6992
	err := log.Create(&logInfo)
	if err != nil && err.Error() != "collection already exists" {
		return nil, maybeUnauthorized(err, "cannot create log collection")
	}
	st.runner = txn.NewRunner(db.C("txns"))
	st.runner.ChangeLog(db.C("txns.log"))
	st.watcher = watcher.New(db.C("txns.log"))
	st.pwatcher = presence.NewWatcher(pdb.C("presence"))
	for _, item := range indexes {
		index := mgo.Index{Key: item.key}
		if err := db.C(item.collection).EnsureIndex(index); err != nil {
			return nil, fmt.Errorf("cannot create database index: %v", err)
		}
	}
	st.transactionHooks = make(chan ([]transactionHook), 1)
	st.transactionHooks <- nil
	return st, nil
}

// Addresses returns the list of addresses used to connect to the state.
func (st *State) Addresses() ([]string, error) {
	stateAddrs := st.db.Session.LiveServers()
	if len(stateAddrs) == 0 {
		return nil, stderrors.New("unable to find state addresses")
	}
	return stateAddrs, nil
}

// APIAddresses returns the list of addresses used to connect to the API.
func (st *State) APIAddresses() ([]string, error) {
	stateAddrs, err := st.Addresses()
	if err != nil {
		return nil, err
	}
	config, err := st.EnvironConfig()
	if err != nil {
		return nil, err
	}
	apiAddrs := make([]string, 0, len(stateAddrs))
	apiPortSuffix := fmt.Sprintf(":%d", config.APIPort())
	for _, stateAddr := range stateAddrs {
		i := strings.LastIndex(stateAddr, ":")
		apiAddrs = append(apiAddrs, stateAddr[:i]+apiPortSuffix)
	}
	return apiAddrs, nil
}

// CACert returns the certificate used to validate the state connection.
func (st *State) CACert() (cert []byte) {
	return append(cert, st.info.CACert...)
}

func (st *State) Close() error {
	err1 := st.watcher.Stop()
	err2 := st.pwatcher.Stop()
	st.mu.Lock()
	var err3 error
	if st.allManager != nil {
		err3 = st.allManager.Stop()
	}
	st.mu.Unlock()
	st.db.Session.Close()
	for _, err := range []error{err1, err2, err3} {
		if err != nil {
			return err
		}
	}
	return nil
}
<commit_msg>Revert a change that was not supposed to be committed<commit_after>// Copyright 2012, 2013 Canonical Ltd.
// Licensed under the AGPLv3, see LICENCE file for details.

package state

import (
	"crypto/tls"
	"crypto/x509"
	stderrors "errors"
	"fmt"
	"net"
	"strings"
	"time"

	"labix.org/v2/mgo"
	"labix.org/v2/mgo/txn"

	"launchpad.net/juju-core/cert"
	"launchpad.net/juju-core/constraints"
	"launchpad.net/juju-core/environs/config"
	"launchpad.net/juju-core/errors"
	"launchpad.net/juju-core/state/presence"
	"launchpad.net/juju-core/state/watcher"
	"launchpad.net/juju-core/utils"
)

// Info encapsulates information about cluster of
// servers holding juju state and can be used to make a
// connection to that cluster.
type Info struct {
	// Addrs gives the addresses of the MongoDB servers for the state.
	// Each address should be in the form address:port.
	Addrs []string

	// CACert holds the CA certificate that will be used
	// to validate the state server's certificate, in PEM format.
	CACert []byte

	// Tag holds the name of the entity that is connecting.
	// It should be empty when connecting as an administrator.
	Tag string

	// Password holds the password for the connecting entity.
	Password string
}

// DialOpts holds configuration parameters that control the
// Dialing behavior when connecting to a state server.
type DialOpts struct {
	// Timeout is the amount of time to wait contacting
	// a state server.
	Timeout time.Duration
}

// DefaultDialOpts returns a DialOpts representing the default
// parameters for contacting a state server.
func DefaultDialOpts() DialOpts {
	return DialOpts{
		Timeout: 10 * time.Minute,
	}
}

// Open connects to the server described by the given
// info, waits for it to be initialized, and returns a new State
// representing the environment connected to.
// It returns unauthorizedError if access is unauthorized.
func Open(info *Info, opts DialOpts) (*State, error) {
	logger.Infof("opening state; mongo addresses: %q; entity %q", info.Addrs, info.Tag)
	if len(info.Addrs) == 0 {
		return nil, stderrors.New("no mongo addresses")
	}
	if len(info.CACert) == 0 {
		return nil, stderrors.New("missing CA certificate")
	}
	xcert, err := cert.ParseCert(info.CACert)
	if err != nil {
		return nil, fmt.Errorf("cannot parse CA certificate: %v", err)
	}
	pool := x509.NewCertPool()
	pool.AddCert(xcert)
	tlsConfig := &tls.Config{
		RootCAs:    pool,
		ServerName: "anything",
	}
	dial := func(addr net.Addr) (net.Conn, error) {
		c, err := net.Dial("tcp", addr.String())
		if err != nil {
			logger.Debugf("connection failed, will retry: %v", err)
			return nil, err
		}
		cc := tls.Client(c, tlsConfig)
		if err := cc.Handshake(); err != nil {
			logger.Errorf("TLS handshake failed: %v", err)
			return nil, err
		}
		return cc, nil
	}
	session, err := mgo.DialWithInfo(&mgo.DialInfo{
		Addrs:   info.Addrs,
		Timeout: opts.Timeout,
		Dial:    dial,
	})
	if err != nil {
		return nil, err
	}
	logger.Infof("connection established")
	st, err := newState(session, info)
	if err != nil {
		session.Close()
		return nil, err
	}
	return st, nil
}

// Initialize sets up an initial empty state and returns it.
// This needs to be performed only once for a given environment.
// It returns unauthorizedError if access is unauthorized.
func Initialize(info *Info, cfg *config.Config, opts DialOpts) (rst *State, err error) {
	st, err := Open(info, opts)
	if err != nil {
		return nil, err
	}
	defer func() {
		if err != nil {
			st.Close()
		}
	}()
	// A valid environment is used as a signal that the
	// state has already been initalized. If this is the case
	// do nothing.
	if _, err := st.Environment(); err == nil {
		return st, nil
	} else if !errors.IsNotFoundError(err) {
		return nil, err
	}
	logger.Infof("initializing environment")
	if err := checkEnvironConfig(cfg); err != nil {
		return nil, err
	}
	uuid, err := utils.NewUUID()
	if err != nil {
		return nil, fmt.Errorf("environment UUID cannot be created: %v", err)
	}
	ops := []txn.Op{
		createConstraintsOp(st, environGlobalKey, constraints.Value{}),
		createSettingsOp(st, environGlobalKey, cfg.AllAttrs()),
		createEnvironmentOp(st, cfg.Name(), uuid.String()),
	}
	if err := st.runTransaction(ops); err == txn.ErrAborted {
		// The config was created in the meantime.
		return st, nil
	} else if err != nil {
		return nil, err
	}
	return st, nil
}

var indexes = []struct {
	collection string
	key        []string
}{
	// After the first public release, do not remove entries from here
	// without adding them to a list of indexes to drop, to ensure
	// old databases are modified to have the correct indexes.
	{"relations", []string{"endpoints.relationname"}},
	{"relations", []string{"endpoints.servicename"}},
	{"units", []string{"service"}},
	{"units", []string{"principal"}},
	{"units", []string{"machineid"}},
	{"users", []string{"name"}},
}

// The capped collection used for transaction logs defaults to 10MB.
// It's tweaked in export_test.go to 1MB to avoid the overhead of
// creating and deleting the large file repeatedly in tests.
var (
	logSize      = 10000000
	logSizeTests = 1000000
)

func maybeUnauthorized(err error, msg string) error {
	if err == nil {
		return nil
	}
	// Unauthorized access errors have no error code,
	// just a simple error string.
	if err.Error() == "auth fails" {
		return errors.NewUnauthorizedError(err, msg)
	}
	if err, ok := err.(*mgo.QueryError); ok && err.Code == 10057 {
		return errors.NewUnauthorizedError(err, msg)
	}
	return fmt.Errorf("%s: %v", msg, err)
}

func newState(session *mgo.Session, info *Info) (*State, error) {
	db := session.DB("juju")
	pdb := session.DB("presence")
	if info.Tag != "" {
		if err := db.Login(info.Tag, info.Password); err != nil {
			return nil, maybeUnauthorized(err, fmt.Sprintf("cannot log in to juju database as %q", info.Tag))
		}
		if err := pdb.Login(info.Tag, info.Password); err != nil {
			return nil, maybeUnauthorized(err, fmt.Sprintf("cannot log in to presence database as %q", info.Tag))
		}
	} else if info.Password != "" {
		admin := session.DB("admin")
		if err := admin.Login("admin", info.Password); err != nil {
			return nil, maybeUnauthorized(err, "cannot log in to admin database")
		}
	}
	st := &State{
		info:           info,
		db:             db,
		environments:   db.C("environments"),
		charms:         db.C("charms"),
		machines:       db.C("machines"),
		containerRefs:  db.C("containerRefs"),
		instanceData:   db.C("instanceData"),
		relations:      db.C("relations"),
		relationScopes: db.C("relationscopes"),
		services:       db.C("services"),
		minUnits:       db.C("minunits"),
		settings:       db.C("settings"),
		settingsrefs:   db.C("settingsrefs"),
		constraints:    db.C("constraints"),
		units:          db.C("units"),
		users:          db.C("users"),
		presence:       pdb.C("presence"),
		cleanups:       db.C("cleanups"),
		annotations:    db.C("annotations"),
		statuses:       db.C("statuses"),
	}
	log := db.C("txns.log")
	logInfo := mgo.CollectionInfo{Capped: true, MaxBytes: logSize}
	// The lack of error code for this error was reported upstream:
	//     https://jira.klmongodb.org/browse/SERVER-6992
	err := log.Create(&logInfo)
	if err != nil && err.Error() != "collection already exists" {
		return nil, maybeUnauthorized(err, "cannot create log collection")
	}
	st.runner = txn.NewRunner(db.C("txns"))
	st.runner.ChangeLog(db.C("txns.log"))
	st.watcher = watcher.New(db.C("txns.log"))
	st.pwatcher = presence.NewWatcher(pdb.C("presence"))
	for _, item := range indexes {
		index := mgo.Index{Key: item.key}
		if err := db.C(item.collection).EnsureIndex(index); err != nil {
			return nil, fmt.Errorf("cannot create database index: %v", err)
		}
	}
	st.transactionHooks = make(chan ([]transactionHook), 1)
	st.transactionHooks <- nil
	return st, nil
}

// Addresses returns the list of addresses used to connect to the state.
func (st *State) Addresses() ([]string, error) {
	stateAddrs := st.db.Session.LiveServers()
	if len(stateAddrs) == 0 {
		return nil, stderrors.New("unable to find state addresses")
	}
	return stateAddrs, nil
}

// APIAddresses returns the list of addresses used to connect to the API.
func (st *State) APIAddresses() ([]string, error) {
	stateAddrs, err := st.Addresses()
	if err != nil {
		return nil, err
	}
	config, err := st.EnvironConfig()
	if err != nil {
		return nil, err
	}
	apiAddrs := make([]string, 0, len(stateAddrs))
	apiPortSuffix := fmt.Sprintf(":%d", config.APIPort())
	for _, stateAddr := range stateAddrs {
		i := strings.LastIndex(stateAddr, ":")
		apiAddrs = append(apiAddrs, stateAddr[:i]+apiPortSuffix)
	}
	return apiAddrs, nil
}

// CACert returns the certificate used to validate the state connection.
func (st *State) CACert() (cert []byte) {
	return append(cert, st.info.CACert...)
}

func (st *State) Close() error {
	err1 := st.watcher.Stop()
	err2 := st.pwatcher.Stop()
	st.mu.Lock()
	var err3 error
	if st.allManager != nil {
		err3 = st.allManager.Stop()
	}
	st.mu.Unlock()
	st.db.Session.Close()
	for _, err := range []error{err1, err2, err3} {
		if err != nil {
			return err
		}
	}
	return nil
}
<|endoftext|>


<commit_before>package php

import (
	"stephensearles.com/php/ast"
	"stephensearles.com/php/token"
)

func (p *Parser) parseStmt() ast.Statement {
	switch p.current.typ {
	case token.BlockBegin:
		p.backup()
		return p.parseBlock()
	case token.Global:
		p.next()
		g := &ast.GlobalDeclaration{
			Identifiers: make([]*ast.Variable, 0, 1),
		}
		for p.current.typ == token.VariableOperator {
			variable, ok := p.parseVariable().(*ast.Variable)
			if !ok {
				p.errorf("global declarations must be of standard variables")
				break
			}
			g.Identifiers = append(g.Identifiers, variable)
			if p.peek().typ != token.Comma {
				break
			}
			p.expect(token.Comma)
			p.next()
		}
		p.expectStmtEnd()
		return g
	case token.Namespace:
		p.expect(token.Identifier)
		p.expectStmtEnd()
		// We are ignoring this for now
		return nil
	case token.Use:
		p.expect(token.Identifier)
		if p.peek().typ == token.AsOperator {
			p.expect(token.AsOperator)
			p.expect(token.Identifier)
		}
		p.expectStmtEnd()
		// We are ignoring this for now
		return nil
	case token.Static:
		s := &ast.StaticVariableDeclaration{Declarations: make([]ast.Expression, 0)}
		for {
			p.expect(token.VariableOperator)
			p.expect(token.Identifier)
			v := ast.NewVariable(p.current.val)
			if p.peek().typ == token.AssignmentOperator {
				p.expect(token.AssignmentOperator)
				op := p.current.val
				p.expect(token.Null, token.StringLiteral, token.BooleanLiteral, token.NumberLiteral, token.Array)
				switch p.current.typ {
				case token.Array:
					s.Declarations = append(s.Declarations, &ast.AssignmentExpression{Assignee: v, Value: p.parseArrayDeclaration(), Operator: op})
				default:
					s.Declarations = append(s.Declarations, &ast.AssignmentExpression{Assignee: v, Value: p.parseLiteral(), Operator: op})
				}
			}
			s.Declarations = append(s.Declarations, v)
			if p.peek().typ != token.Comma {
				break
			}
			p.next()
		}
		p.expectStmtEnd()
		return s
	case token.VariableOperator, token.UnaryOperator:
		expr := ast.ExpressionStmt{p.parseExpression()}
		p.expectStmtEnd()
		return expr
	case token.Print:
		requireParen := false
		if p.peek().typ == token.OpenParen {
			p.expect(token.OpenParen)
			requireParen = true
		}
		stmt := ast.Echo(p.parseNextExpression())
		if requireParen {
			p.expect(token.CloseParen)
		}
		p.expectStmtEnd()
		return stmt
	case token.Function:
		return p.parseFunctionStmt()
	case token.PHPEnd:
		if p.peek().typ == token.EOF {
			return nil
		}
		var expr ast.Statement
		if p.accept(token.HTML) {
			expr = ast.Echo(&ast.Literal{Type: ast.String, Value: p.current.val})
		}
		p.next()
		if p.current.typ != token.EOF {
			p.expectCurrent(token.PHPBegin)
		}
		return expr
	case token.Echo:
		exprs := []ast.Expression{
			p.parseNextExpression(),
		}
		for p.peek().typ == token.Comma {
			p.expect(token.Comma)
			exprs = append(exprs, p.parseNextExpression())
		}
		p.expectStmtEnd()
		return ast.Echo(exprs...)
	case token.If:
		return p.parseIf()
	case token.While:
		return p.parseWhile()
	case token.Do:
		return p.parseDo()
	case token.For:
		return p.parseFor()
	case token.Foreach:
		return p.parseForeach()
	case token.Switch:
		return p.parseSwitch()
	case token.Abstract, token.Final, token.Class:
		return p.parseClass()
	case token.Interface:
		return p.parseInterface()
	case token.Return:
		p.next()
		stmt := ast.ReturnStmt{}
		if p.current.typ != token.StatementEnd {
			stmt.Expression = p.parseExpression()
			p.expectStmtEnd()
		}
		return stmt
	case token.Break:
		p.next()
		stmt := ast.BreakStmt{}
		if p.current.typ != token.StatementEnd {
			stmt.Expression = p.parseExpression()
			p.expectStmtEnd()
		}
		return stmt
	case token.Continue:
		p.next()
		stmt := ast.ContinueStmt{}
		if p.current.typ != token.StatementEnd {
			stmt.Expression = p.parseExpression()
			p.expectStmtEnd()
		}
		return stmt
	case token.Throw:
		stmt := ast.ThrowStmt{Expression: p.parseNextExpression()}
		p.expectStmtEnd()
		return stmt
	case token.Exit:
		stmt := ast.ExitStmt{}
		if p.peek().typ == token.OpenParen {
			p.expect(token.OpenParen)
			if p.peek().typ != token.CloseParen {
				stmt.Expression = p.parseNextExpression()
			}
			p.expect(token.CloseParen)
		}
		p.expectStmtEnd()
		return stmt
	case token.Try:
		stmt := &ast.TryStmt{}
		stmt.TryBlock = p.parseBlock()
		for p.expect(token.Catch); p.current.typ == token.Catch; p.next() {
			caught := &ast.CatchStmt{}
			p.expect(token.OpenParen)
			p.expect(token.Identifier)
			caught.CatchType = p.current.val
			p.expect(token.VariableOperator)
			p.expect(token.Identifier)
			caught.CatchVar = ast.NewVariable(p.current.val)
			p.expect(token.CloseParen)
			caught.CatchBlock = p.parseBlock()
			stmt.CatchStmts = append(stmt.CatchStmts, caught)
		}
		p.backup()
		return stmt
	case token.IgnoreErrorOperator:
		// Ignore this operator
		p.next()
		return p.parseStmt()
	case token.StatementEnd:
		// this is an empty statement
		return &ast.EmptyStatement{}
	default:
		expr := p.parseExpression()
		if expr != nil {
			p.expectStmtEnd()
			return ast.ExpressionStmt{expr}
		}
		p.errorf("Found %s, statement or expression", p.current)
		return nil
	}
}

func (p *Parser) expectStmtEnd() {
	if p.peek().typ != token.PHPEnd {
		p.expect(token.StatementEnd)
	}
}
<commit_msg>Fixing parsing static keyworkd, checking if it should be an expression<commit_after>package php

import (
	"stephensearles.com/php/ast"
	"stephensearles.com/php/token"
)

func (p *Parser) parseStmt() ast.Statement {
	switch p.current.typ {
	case token.BlockBegin:
		p.backup()
		return p.parseBlock()
	case token.Global:
		p.next()
		g := &ast.GlobalDeclaration{
			Identifiers: make([]*ast.Variable, 0, 1),
		}
		for p.current.typ == token.VariableOperator {
			variable, ok := p.parseVariable().(*ast.Variable)
			if !ok {
				p.errorf("global declarations must be of standard variables")
				break
			}
			g.Identifiers = append(g.Identifiers, variable)
			if p.peek().typ != token.Comma {
				break
			}
			p.expect(token.Comma)
			p.next()
		}
		p.expectStmtEnd()
		return g
	case token.Namespace:
		p.expect(token.Identifier)
		p.expectStmtEnd()
		// We are ignoring this for now
		return nil
	case token.Use:
		p.expect(token.Identifier)
		if p.peek().typ == token.AsOperator {
			p.expect(token.AsOperator)
			p.expect(token.Identifier)
		}
		p.expectStmtEnd()
		// We are ignoring this for now
		return nil
	case token.Static:
		if p.peek().typ == token.ScopeResolutionOperator {
			expr := p.parseExpression()
			p.expectStmtEnd()
			return expr
		}
		s := &ast.StaticVariableDeclaration{Declarations: make([]ast.Expression, 0)}
		for {
			p.expect(token.VariableOperator)
			p.expect(token.Identifier)
			v := ast.NewVariable(p.current.val)
			if p.peek().typ == token.AssignmentOperator {
				p.expect(token.AssignmentOperator)
				op := p.current.val
				p.expect(token.Null, token.StringLiteral, token.BooleanLiteral, token.NumberLiteral, token.Array)
				switch p.current.typ {
				case token.Array:
					s.Declarations = append(s.Declarations, &ast.AssignmentExpression{Assignee: v, Value: p.parseArrayDeclaration(), Operator: op})
				default:
					s.Declarations = append(s.Declarations, &ast.AssignmentExpression{Assignee: v, Value: p.parseLiteral(), Operator: op})
				}
			}
			s.Declarations = append(s.Declarations, v)
			if p.peek().typ != token.Comma {
				break
			}
			p.next()
		}
		p.expectStmtEnd()
		return s
	case token.VariableOperator, token.UnaryOperator:
		expr := ast.ExpressionStmt{p.parseExpression()}
		p.expectStmtEnd()
		return expr
	case token.Print:
		requireParen := false
		if p.peek().typ == token.OpenParen {
			p.expect(token.OpenParen)
			requireParen = true
		}
		stmt := ast.Echo(p.parseNextExpression())
		if requireParen {
			p.expect(token.CloseParen)
		}
		p.expectStmtEnd()
		return stmt
	case token.Function:
		return p.parseFunctionStmt()
	case token.PHPEnd:
		if p.peek().typ == token.EOF {
			return nil
		}
		var expr ast.Statement
		if p.accept(token.HTML) {
			expr = ast.Echo(&ast.Literal{Type: ast.String, Value: p.current.val})
		}
		p.next()
		if p.current.typ != token.EOF {
			p.expectCurrent(token.PHPBegin)
		}
		return expr
	case token.Echo:
		exprs := []ast.Expression{
			p.parseNextExpression(),
		}
		for p.peek().typ == token.Comma {
			p.expect(token.Comma)
			exprs = append(exprs, p.parseNextExpression())
		}
		p.expectStmtEnd()
		return ast.Echo(exprs...)
	case token.If:
		return p.parseIf()
	case token.While:
		return p.parseWhile()
	case token.Do:
		return p.parseDo()
	case token.For:
		return p.parseFor()
	case token.Foreach:
		return p.parseForeach()
	case token.Switch:
		return p.parseSwitch()
	case token.Abstract, token.Final, token.Class:
		return p.parseClass()
	case token.Interface:
		return p.parseInterface()
	case token.Return:
		p.next()
		stmt := ast.ReturnStmt{}
		if p.current.typ != token.StatementEnd {
			stmt.Expression = p.parseExpression()
			p.expectStmtEnd()
		}
		return stmt
	case token.Break:
		p.next()
		stmt := ast.BreakStmt{}
		if p.current.typ != token.StatementEnd {
			stmt.Expression = p.parseExpression()
			p.expectStmtEnd()
		}
		return stmt
	case token.Continue:
		p.next()
		stmt := ast.ContinueStmt{}
		if p.current.typ != token.StatementEnd {
			stmt.Expression = p.parseExpression()
			p.expectStmtEnd()
		}
		return stmt
	case token.Throw:
		stmt := ast.ThrowStmt{Expression: p.parseNextExpression()}
		p.expectStmtEnd()
		return stmt
	case token.Exit:
		stmt := ast.ExitStmt{}
		if p.peek().typ == token.OpenParen {
			p.expect(token.OpenParen)
			if p.peek().typ != token.CloseParen {
				stmt.Expression = p.parseNextExpression()
			}
			p.expect(token.CloseParen)
		}
		p.expectStmtEnd()
		return stmt
	case token.Try:
		stmt := &ast.TryStmt{}
		stmt.TryBlock = p.parseBlock()
		for p.expect(token.Catch); p.current.typ == token.Catch; p.next() {
			caught := &ast.CatchStmt{}
			p.expect(token.OpenParen)
			p.expect(token.Identifier)
			caught.CatchType = p.current.val
			p.expect(token.VariableOperator)
			p.expect(token.Identifier)
			caught.CatchVar = ast.NewVariable(p.current.val)
			p.expect(token.CloseParen)
			caught.CatchBlock = p.parseBlock()
			stmt.CatchStmts = append(stmt.CatchStmts, caught)
		}
		p.backup()
		return stmt
	case token.IgnoreErrorOperator:
		// Ignore this operator
		p.next()
		return p.parseStmt()
	case token.StatementEnd:
		// this is an empty statement
		return &ast.EmptyStatement{}
	default:
		expr := p.parseExpression()
		if expr != nil {
			p.expectStmtEnd()
			return ast.ExpressionStmt{expr}
		}
		p.errorf("Found %s, statement or expression", p.current)
		return nil
	}
}

func (p *Parser) expectStmtEnd() {
	if p.peek().typ != token.PHPEnd {
		p.expect(token.StatementEnd)
	}
}
<|endoftext|>


<commit_before>package main

import (
	"encoding/json"
	"sync"
	"time"
)

// Statistics keeps stats about the current operation of the program. It is
// meant to keep snapshot-in-time stats, as opposed to counters or timers that
// statsd offers.
//
// Statistics may be exposed by APIs that allow human- or machine-readable
// monitoring.
type Statistics struct {
	files map[string]*FileStatistics

	// Synchronizes access to the Files map
	filesLock sync.RWMutex
}

const (
	// The status of the file has not yet been explicitly set.
	fileStatusUnknown = "unknown"

	// The file is currently being read.
	fileStatusReading = "reading"

	// The file has been read to the end. In a few minutes, the file will be
	// closed. Or, if more data is written, the status will go back to reading.
	fileStatusEof = "eof"

	// The file is no longer being read. The file has been read to EOF and it
	// has not yet been reopened. If the file has been deleted, it will never
	// be reopened and will remain in this status until the process restarts.
	fileStatusClosed = "closed"
)

type FileStatistics struct {
	Status string `json:"status"`

	// The current position (in bytes) that has been read into the file. This
	// might be greater than SnapshotPosition if there are lines buffered into
	// memory that haven't been acknowledged by the server
	Position int64 `json:"position"`

	// The last time the file was read from into the in-memory buffer.
	LastRead time.Time `json:"last_read"`

	// The current position (in bytes) that has been successfully sent and
	// acknowledged by the remote server.
	SnapshotPosition int64 `json:"snapshot_position"`

	// The last time a line from this file was successfully sent and acknowledged
	// by the remote server.
	LastSnapshot time.Time `json:"last_snapshot"`
}

var GlobalStatistics *Statistics = NewStatistics()

func NewStatistics() *Statistics {
	return &Statistics{}
}

func (s *Statistics) SetFileStatus(filePath string, status string) {
	s.ensureFileStatisticsCreated(filePath)

	stats := s.GetFileStatistics(filePath)
	stats.Status = status
}

func (s *Statistics) SetFilePosition(filePath string, position int64) {
	s.ensureFileStatisticsCreated(filePath)

	stats := s.GetFileStatistics(filePath)
	stats.Position = position
	stats.LastRead = time.Now()
}

func (s *Statistics) SetFileSnapshotPosition(filePath string, snapshotPosition int64) {
	s.ensureFileStatisticsCreated(filePath)

	stats := s.GetFileStatistics(filePath)
	stats.SnapshotPosition = snapshotPosition
	stats.LastSnapshot = time.Now()
}

func (s *Statistics) GetFileStatistics(filePath string) *FileStatistics {
	s.filesLock.RLock()
	defer s.filesLock.RUnlock()

	return s.files[filePath]
}

func (s *Statistics) ensureFileStatisticsCreated(filePath string) {
	// Fast check
	if s.files == nil {
		s.filesLock.Lock()
		// Check again in the critical region
		if s.files == nil {
			s.files = make(map[string]*FileStatistics)
		}
		s.filesLock.Unlock()
	}

	// Fast check
	if _, ok := s.files[filePath]; !ok {
		s.filesLock.Lock()
		// Check again in the critical region
		if _, ok := s.files[filePath]; !ok {
			s.files[filePath] = &FileStatistics{}
		}
		s.filesLock.Unlock()
	}
}

func (s *Statistics) MarshalJSON() ([]byte, error) {
	structure := map[string]interface{}{
		"files": s.files,
	}

	return json.Marshal(structure)
}
<commit_msg>Constructor handles this initialization<commit_after>package main

import (
	"encoding/json"
	"sync"
	"time"
)

// Statistics keeps stats about the current operation of the program. It is
// meant to keep snapshot-in-time stats, as opposed to counters or timers that
// statsd offers.
//
// Statistics may be exposed by APIs that allow human- or machine-readable
// monitoring.
type Statistics struct {
	files map[string]*FileStatistics

	// Synchronizes access to the Files map
	filesLock sync.RWMutex
}

const (
	// The status of the file has not yet been explicitly set.
	fileStatusUnknown = "unknown"

	// The file is currently being read.
	fileStatusReading = "reading"

	// The file has been read to the end. In a few minutes, the file will be
	// closed. Or, if more data is written, the status will go back to reading.
	fileStatusEof = "eof"

	// The file is no longer being read. The file has been read to EOF and it
	// has not yet been reopened. If the file has been deleted, it will never
	// be reopened and will remain in this status until the process restarts.
	fileStatusClosed = "closed"
)

type FileStatistics struct {
	Status string `json:"status"`

	// The current position (in bytes) that has been read into the file. This
	// might be greater than SnapshotPosition if there are lines buffered into
	// memory that haven't been acknowledged by the server
	Position int64 `json:"position"`

	// The last time the file was read from into the in-memory buffer.
	LastRead time.Time `json:"last_read"`

	// The current position (in bytes) that has been successfully sent and
	// acknowledged by the remote server.
	SnapshotPosition int64 `json:"snapshot_position"`

	// The last time a line from this file was successfully sent and acknowledged
	// by the remote server.
	LastSnapshot time.Time `json:"last_snapshot"`
}

var GlobalStatistics *Statistics = NewStatistics()

func NewStatistics() *Statistics {
	return &Statistics{
		files: make(map[string]*FileStatistics),
	}
}

func (s *Statistics) SetFileStatus(filePath string, status string) {
	s.ensureFileStatisticsCreated(filePath)

	stats := s.GetFileStatistics(filePath)
	stats.Status = status
}

func (s *Statistics) SetFilePosition(filePath string, position int64) {
	s.ensureFileStatisticsCreated(filePath)

	stats := s.GetFileStatistics(filePath)
	stats.Position = position
	stats.LastRead = time.Now()
}

func (s *Statistics) SetFileSnapshotPosition(filePath string, snapshotPosition int64) {
	s.ensureFileStatisticsCreated(filePath)

	stats := s.GetFileStatistics(filePath)
	stats.SnapshotPosition = snapshotPosition
	stats.LastSnapshot = time.Now()
}

func (s *Statistics) GetFileStatistics(filePath string) *FileStatistics {
	s.filesLock.RLock()
	defer s.filesLock.RUnlock()

	return s.files[filePath]
}

func (s *Statistics) ensureFileStatisticsCreated(filePath string) {
	// Fast check
	if _, ok := s.files[filePath]; !ok {
		s.filesLock.Lock()
		// Check again in the critical region
		if _, ok := s.files[filePath]; !ok {
			s.files[filePath] = &FileStatistics{}
		}
		s.filesLock.Unlock()
	}
}

func (s *Statistics) MarshalJSON() ([]byte, error) {
	structure := map[string]interface{}{
		"files": s.files,
	}

	return json.Marshal(structure)
}
<|endoftext|>