<commit_before>package minion

import (
	"bytes"
	"encoding/json"
	"log"
	"os/exec"
	"path/filepath"
	"strconv"
	"strings"
	"time"

	"github.com/dnaeon/gru/classifier"
	"github.com/dnaeon/gru/task"
	"github.com/dnaeon/gru/utils"

	"code.google.com/p/go-uuid/uuid"
	etcdclient "github.com/coreos/etcd/client"
	"golang.org/x/net/context"
)

// Minions keyspace in etcd
const EtcdMinionSpace = "/gru/minion"

// Etcd Minion
type etcdMinion struct {
	// Name of the minion
	name string

	// Minion root directory in etcd
	rootDir string

	// Minion queue directory in etcd
	queueDir string

	// Log directory of previously executed tasks
	logDir string

	// Classifier directory in etcd
	classifierDir string

	// Minion unique identifier
	id uuid.UUID

	// KeysAPI client to etcd
	kapi etcdclient.KeysAPI

	// Channel over which tasks are sent for processing
	taskQueue chan *task.Task
}

// Creates a new etcd minion
func NewEtcdMinion(name string, cfg etcdclient.Config) Minion {
	c, err := etcdclient.New(cfg)
	if err != nil {
		log.Fatal(err)
	}

	kapi := etcdclient.NewKeysAPI(c)
	id := utils.GenerateUUID(name)
	rootDir := filepath.Join(EtcdMinionSpace, id.String())
	queueDir := filepath.Join(rootDir, "queue")
	classifierDir := filepath.Join(rootDir, "classifier")
	logDir := filepath.Join(rootDir, "log")
	taskQueue := make(chan *task.Task)

	m := &etcdMinion{
		name:          name,
		rootDir:       rootDir,
		queueDir:      queueDir,
		classifierDir: classifierDir,
		logDir:        logDir,
		id:            id,
		kapi:          kapi,
		taskQueue:     taskQueue,
	}

	return m
}

// Set the human-readable name of the minion in etcd
func (m *etcdMinion) setName() error {
	nameKey := filepath.Join(m.rootDir, "name")
	opts := &etcdclient.SetOptions{
		PrevExist: etcdclient.PrevIgnore,
	}

	_, err := m.kapi.Set(context.Background(), nameKey, m.Name(), opts)
	if err != nil {
		log.Printf("Failed to set name of minion: %s\n", err)
	}

	return err
}

// Set the time the minion was last seen in seconds since the Epoch
func (m *etcdMinion) setLastseen(s int64) error {
	lastseenKey := filepath.Join(m.rootDir, "lastseen")
	lastseenValue := strconv.FormatInt(s, 10)
	opts := &etcdclient.SetOptions{
		PrevExist: etcdclient.PrevIgnore,
	}

	_, err := m.kapi.Set(context.Background(), lastseenKey, lastseenValue, opts)
	if err != nil {
		log.Printf("Failed to set lastseen time: %s\n", err)
	}

	return err
}

// Checks for any pending tasks and sends them
// for processing if there are any
func (m *etcdMinion) checkQueue() error {
	opts := &etcdclient.GetOptions{
		Recursive: true,
		Sort:      true,
	}

	// Get backlog tasks if any
	resp, err := m.kapi.Get(context.Background(), m.queueDir, opts)
	if err != nil {
		log.Printf("Failed to get backlog tasks: %s\n", err)
		return err
	}

	backlog := resp.Node.Nodes
	if len(backlog) == 0 {
		// No backlog tasks found
		return nil
	}

	log.Printf("Found %d pending tasks in queue", len(backlog))
	for _, node := range backlog {
		t, err := EtcdUnmarshalTask(node)
		m.kapi.Delete(context.Background(), node.Key, nil)

		if err != nil {
			continue
		}

		m.taskQueue <- t
	}

	return nil
}

// Runs periodic jobs such as refreshing classifiers and
// updating the lastseen time every fifteen minutes
func (m *etcdMinion) periodicRunner() {
	schedule := time.Minute * 15
	ticker := time.NewTicker(schedule)
	log.Printf("Periodic scheduler set to run every %s\n", schedule)

	for now := range ticker.C {
		// Run any periodic jobs
		m.Classify()
		m.setLastseen(now.Unix())
	}
}

// Processes new tasks
func (m *etcdMinion) processTask(t *task.Task) error {
	var buf bytes.Buffer

	// Update state of task to indicate that we are now processing it
	t.State = task.TaskStateProcessing
	m.SaveTaskResult(t)

	cmd := exec.Command(t.Command, t.Args...)
	cmd.Stdout = &buf
	cmd.Stderr = &buf

	log.Printf("Processing task %s\n", t.TaskID)

	cmdError := cmd.Run()
	t.TimeProcessed = time.Now().Unix()
	t.Result = buf.String()

	if cmdError != nil {
		log.Printf("Failed to process task %s\n", t.TaskID)
		t.Error = cmdError.Error()
		t.State = task.TaskStateFailed
	} else {
		log.Printf("Finished processing task %s\n", t.TaskID)
		t.State = task.TaskStateSuccess
	}

	m.SaveTaskResult(t)

	return cmdError
}

// Saves a task in etcd
func (m *etcdMinion) SaveTaskResult(t *task.Task) error {
	taskKey := filepath.Join(m.logDir, t.TaskID.String())

	data, err := json.Marshal(t)
	if err != nil {
		log.Printf("Failed to serialize task %s: %s\n", t.TaskID, err)
		return err
	}

	opts := &etcdclient.SetOptions{
		PrevExist: etcdclient.PrevIgnore,
	}

	_, err = m.kapi.Set(context.Background(), taskKey, string(data), opts)
	if err != nil {
		log.Printf("Failed to save task %s: %s\n", t.TaskID, err)
	}

	return err
}

// Unmarshals task from etcd
func EtcdUnmarshalTask(node *etcdclient.Node) (*task.Task, error) {
	task := new(task.Task)
	err := json.Unmarshal([]byte(node.Value), &task)

	return task, err
}

// Returns the minion unique identifier
func (m *etcdMinion) ID() uuid.UUID {
	return m.id
}

// Returns the assigned name of the minion
func (m *etcdMinion) Name() string {
	return m.name
}

// Classifies the minion
func (m *etcdMinion) Classify() error {
	// Classifiers in etcd expire after an hour
	opts := &etcdclient.SetOptions{
		PrevExist: etcdclient.PrevIgnore,
		TTL:       time.Hour,
	}

	// Set/update classifiers in etcd
	for key, _ := range classifier.Registry {
		klassifier, err := classifier.Get(key)

		if err != nil {
			log.Printf("Failed to get classifier %s: %s\n", key, err)
			continue
		}

		// Serialize classifier to JSON and save it in etcd
		data, err := json.Marshal(klassifier)
		if err != nil {
			log.Printf("Failed to serialize classifier: %s\n", key)
			continue
		}

		// Classifier key in etcd
		klassifierKey := filepath.Join(m.classifierDir, key)
		_, err = m.kapi.Set(context.Background(), klassifierKey, string(data), opts)

		if err != nil {
			log.Printf("Failed to set classifier %s: %s\n", key, err)
		}
	}

	return nil
}

// Monitors etcd for new tasks
func (m *etcdMinion) TaskListener(c chan<- *task.Task) error {
	log.Printf("Task listener is watching %s\n", m.queueDir)

	watcherOpts := &etcdclient.WatcherOptions{
		Recursive: true,
	}
	watcher := m.kapi.Watcher(m.queueDir, watcherOpts)

	for {
		resp, err := watcher.Next(context.Background())
		if err != nil {
			log.Printf("Failed to receive task: %s\n", err)
			continue
		}

		// Ignore "delete" events when removing a task from the queue
		action := strings.ToLower(resp.Action)
		if strings.EqualFold(action, "delete") {
			continue
		}

		// Unmarshal and remove task from the queue
		t, err := EtcdUnmarshalTask(resp.Node)
		m.kapi.Delete(context.Background(), resp.Node.Key, nil)

		if err != nil {
			log.Printf("Received invalid task %s: %s\n", resp.Node.Key, err)
			continue
		}

		// Update task state and send it for processing
		log.Printf("Received task %s\n", t.TaskID)
		t.State = task.TaskStateQueued
		t.TimeReceived = time.Now().Unix()
		m.SaveTaskResult(t)
		c <- t
	}

	return nil
}

// Processes new tasks
func (m *etcdMinion) TaskRunner(c <-chan *task.Task) error {
	log.Println("Starting task runner")

	for t := range c {
		if t.IsConcurrent {
			go m.processTask(t)
		} else {
			m.processTask(t)
		}
	}

	return nil
}

// Main entry point of the minion
func (m *etcdMinion) Serve() error {
	err := m.setName()
	if err != nil {
		return err
	}

	// Start minion services
	go m.periodicRunner()
	go m.checkQueue()
	go m.TaskRunner(m.taskQueue)
	go m.TaskListener(m.taskQueue)

	log.Printf("Minion %s is ready to serve", m.ID())

	return nil
}

// Stops the minion and performs any cleanup tasks
func (m *etcdMinion) Stop() error {
	log.Println("Minion is shutting down")

	close(m.taskQueue)

	return nil
}
<commit_msg>Don't log an error if there are no pending tasks for processing<commit_after>package minion

import (
	"bytes"
	"encoding/json"
	"log"
	"os/exec"
	"path/filepath"
	"strconv"
	"strings"
	"time"

	"github.com/dnaeon/gru/classifier"
	"github.com/dnaeon/gru/task"
	"github.com/dnaeon/gru/utils"

	"code.google.com/p/go-uuid/uuid"
	etcdclient "github.com/coreos/etcd/client"
	"golang.org/x/net/context"
)

// Minions keyspace in etcd
const EtcdMinionSpace = "/gru/minion"

// Etcd Minion
type etcdMinion struct {
	// Name of the minion
	name string

	// Minion root directory in etcd
	rootDir string

	// Minion queue directory in etcd
	queueDir string

	// Log directory of previously executed tasks
	logDir string

	// Classifier directory in etcd
	classifierDir string

	// Minion unique identifier
	id uuid.UUID

	// KeysAPI client to etcd
	kapi etcdclient.KeysAPI

	// Channel over which tasks are sent for processing
	taskQueue chan *task.Task
}

// Creates a new etcd minion
func NewEtcdMinion(name string, cfg etcdclient.Config) Minion {
	c, err := etcdclient.New(cfg)
	if err != nil {
		log.Fatal(err)
	}

	kapi := etcdclient.NewKeysAPI(c)
	id := utils.GenerateUUID(name)
	rootDir := filepath.Join(EtcdMinionSpace, id.String())
	queueDir := filepath.Join(rootDir, "queue")
	classifierDir := filepath.Join(rootDir, "classifier")
	logDir := filepath.Join(rootDir, "log")
	taskQueue := make(chan *task.Task)

	m := &etcdMinion{
		name:          name,
		rootDir:       rootDir,
		queueDir:      queueDir,
		classifierDir: classifierDir,
		logDir:        logDir,
		id:            id,
		kapi:          kapi,
		taskQueue:     taskQueue,
	}

	return m
}

// Set the human-readable name of the minion in etcd
func (m *etcdMinion) setName() error {
	nameKey := filepath.Join(m.rootDir, "name")
	opts := &etcdclient.SetOptions{
		PrevExist: etcdclient.PrevIgnore,
	}

	_, err := m.kapi.Set(context.Background(), nameKey, m.Name(), opts)
	if err != nil {
		log.Printf("Failed to set name of minion: %s\n", err)
	}

	return err
}

// Set the time the minion was last seen in seconds since the Epoch
func (m *etcdMinion) setLastseen(s int64) error {
	lastseenKey := filepath.Join(m.rootDir, "lastseen")
	lastseenValue := strconv.FormatInt(s, 10)
	opts := &etcdclient.SetOptions{
		PrevExist: etcdclient.PrevIgnore,
	}

	_, err := m.kapi.Set(context.Background(), lastseenKey, lastseenValue, opts)
	if err != nil {
		log.Printf("Failed to set lastseen time: %s\n", err)
	}

	return err
}

// Checks for any pending tasks and sends them
// for processing if there are any
func (m *etcdMinion) checkQueue() error {
	opts := &etcdclient.GetOptions{
		Recursive: true,
		Sort:      true,
	}

	// Get backlog tasks if any
	// If the directory key in etcd is missing that is okay, since
	// it means there are no pending tasks for processing
	resp, err := m.kapi.Get(context.Background(), m.queueDir, opts)
	if err != nil {
		if eerr, ok := err.(etcdclient.Error); !ok || eerr.Code != etcdclient.ErrorCodeKeyNotFound {
			return err
		}
	}

	backlog := resp.Node.Nodes
	if len(backlog) == 0 {
		// No backlog tasks found
		return nil
	}

	log.Printf("Found %d pending tasks in queue", len(backlog))
	for _, node := range backlog {
		t, err := EtcdUnmarshalTask(node)
		m.kapi.Delete(context.Background(), node.Key, nil)

		if err != nil {
			continue
		}

		m.taskQueue <- t
	}

	return nil
}

// Runs periodic jobs such as refreshing classifiers and
// updating the lastseen time every fifteen minutes
func (m *etcdMinion) periodicRunner() {
	schedule := time.Minute * 15
	ticker := time.NewTicker(schedule)
	log.Printf("Periodic scheduler set to run every %s\n", schedule)

	for now := range ticker.C {
		// Run any periodic jobs
		m.Classify()
		m.setLastseen(now.Unix())
	}
}

// Processes new tasks
func (m *etcdMinion) processTask(t *task.Task) error {
	var buf bytes.Buffer

	// Update state of task to indicate that we are now processing it
	t.State = task.TaskStateProcessing
	m.SaveTaskResult(t)

	cmd := exec.Command(t.Command, t.Args...)
	cmd.Stdout = &buf
	cmd.Stderr = &buf

	log.Printf("Processing task %s\n", t.TaskID)

	cmdError := cmd.Run()
	t.TimeProcessed = time.Now().Unix()
	t.Result = buf.String()

	if cmdError != nil {
		log.Printf("Failed to process task %s\n", t.TaskID)
		t.Error = cmdError.Error()
		t.State = task.TaskStateFailed
	} else {
		log.Printf("Finished processing task %s\n", t.TaskID)
		t.State = task.TaskStateSuccess
	}

	m.SaveTaskResult(t)

	return cmdError
}

// Saves a task in etcd
func (m *etcdMinion) SaveTaskResult(t *task.Task) error {
	taskKey := filepath.Join(m.logDir, t.TaskID.String())

	data, err := json.Marshal(t)
	if err != nil {
		log.Printf("Failed to serialize task %s: %s\n", t.TaskID, err)
		return err
	}

	opts := &etcdclient.SetOptions{
		PrevExist: etcdclient.PrevIgnore,
	}

	_, err = m.kapi.Set(context.Background(), taskKey, string(data), opts)
	if err != nil {
		log.Printf("Failed to save task %s: %s\n", t.TaskID, err)
	}

	return err
}

// Unmarshals task from etcd
func EtcdUnmarshalTask(node *etcdclient.Node) (*task.Task, error) {
	task := new(task.Task)
	err := json.Unmarshal([]byte(node.Value), &task)

	return task, err
}

// Returns the minion unique identifier
func (m *etcdMinion) ID() uuid.UUID {
	return m.id
}

// Returns the assigned name of the minion
func (m *etcdMinion) Name() string {
	return m.name
}

// Classifies the minion
func (m *etcdMinion) Classify() error {
	// Classifiers in etcd expire after an hour
	opts := &etcdclient.SetOptions{
		PrevExist: etcdclient.PrevIgnore,
		TTL:       time.Hour,
	}

	// Set/update classifiers in etcd
	for key, _ := range classifier.Registry {
		klassifier, err := classifier.Get(key)

		if err != nil {
			log.Printf("Failed to get classifier %s: %s\n", key, err)
			continue
		}

		// Serialize classifier to JSON and save it in etcd
		data, err := json.Marshal(klassifier)
		if err != nil {
			log.Printf("Failed to serialize classifier: %s\n", key)
			continue
		}

		// Classifier key in etcd
		klassifierKey := filepath.Join(m.classifierDir, key)
		_, err = m.kapi.Set(context.Background(), klassifierKey, string(data), opts)

		if err != nil {
			log.Printf("Failed to set classifier %s: %s\n", key, err)
		}
	}

	return nil
}

// Monitors etcd for new tasks
func (m *etcdMinion) TaskListener(c chan<- *task.Task) error {
	log.Printf("Task listener is watching %s\n", m.queueDir)

	watcherOpts := &etcdclient.WatcherOptions{
		Recursive: true,
	}
	watcher := m.kapi.Watcher(m.queueDir, watcherOpts)

	for {
		resp, err := watcher.Next(context.Background())
		if err != nil {
			log.Printf("Failed to receive task: %s\n", err)
			continue
		}

		// Ignore "delete" events when removing a task from the queue
		action := strings.ToLower(resp.Action)
		if strings.EqualFold(action, "delete") {
			continue
		}

		// Unmarshal and remove task from the queue
		t, err := EtcdUnmarshalTask(resp.Node)
		m.kapi.Delete(context.Background(), resp.Node.Key, nil)

		if err != nil {
			log.Printf("Received invalid task %s: %s\n", resp.Node.Key, err)
			continue
		}

		// Update task state and send it for processing
		log.Printf("Received task %s\n", t.TaskID)
		t.State = task.TaskStateQueued
		t.TimeReceived = time.Now().Unix()
		m.SaveTaskResult(t)
		c <- t
	}

	return nil
}

// Processes new tasks
func (m *etcdMinion) TaskRunner(c <-chan *task.Task) error {
	log.Println("Starting task runner")

	for t := range c {
		if t.IsConcurrent {
			go m.processTask(t)
		} else {
			m.processTask(t)
		}
	}

	return nil
}

// Main entry point of the minion
func (m *etcdMinion) Serve() error {
	err := m.setName()
	if err != nil {
		return err
	}

	// Start minion services
	go m.periodicRunner()
	go m.checkQueue()
	go m.TaskRunner(m.taskQueue)
	go m.TaskListener(m.taskQueue)

	log.Printf("Minion %s is ready to serve", m.ID())

	return nil
}

// Stops the minion and performs any cleanup tasks
func (m *etcdMinion) Stop() error {
	log.Println("Minion is shutting down")

	close(m.taskQueue)

	return nil
}
<|endoftext|>


<commit_before>package main

import (
	"bytes"
	"context"
	"io/ioutil"
	"net"
	"net/http"
	"net/url"
	"os"
	"strconv"
	"sync"
	"sync/atomic"
	"time"

	log "github.com/Sirupsen/logrus"
	"github.com/buger/jsonparser"
	"github.com/gorilla/websocket"
	"github.com/joho/godotenv"
	"github.com/syhlion/requestwork.v2"
	"github.com/urfave/cli"
)

var (
	name     string
	version  string
	cmdStart = cli.Command{
		Name:    "run",
		Aliases: []string{"r"},
		Usage:   "test connect to websocket ",
		Action:  start,
		Flags: []cli.Flag{
			cli.StringFlag{
				Name:  "env-file,e",
				Usage: "import env file",
			},
			cli.BoolFlag{
				Name:  "debug,d",
				Usage: "open debug mode",
			},
		},
	}
	wg        sync.WaitGroup
	listen_wg sync.WaitGroup
)

func start(c *cli.Context) {
	if c.String("env-file") != "" {
		envfile := c.String("env-file")
		//flag.Parse()
		err := godotenv.Load(envfile)
		if err != nil {
			log.Fatal(err)
		}
	}
	ws_api := os.Getenv("GUSHER-CONN-TEST_WS_API")
	if ws_api == "" {
		log.Fatal("empty env GUSHER-CONN-TEST_WS_API")
	}
	ws_auth_api := os.Getenv("GUSHER-CONN-TEST_WS_AUTH_API")
	if ws_auth_api == "" {
		log.Fatal("empty env GUSHER-CONN-TEST_WS_AUTH_API")
	}
	push_api := os.Getenv("GUSHER-CONN-TEST_PUSH_API")
	if push_api == "" {
		log.Fatal("empty env GUSHER-CONN-TEST_PUSH_API")
	}
	jwt := os.Getenv("GUSHER-CONN-TEST_JWT")
	if jwt == "" {
		log.Fatal("empty env GUSHER-CONN-TEST_JWT")
	}
	sub_msg := os.Getenv("GUSHER-CONN-TEST_SUBSCRIBE_MESSAGE")
	if sub_msg == "" {
		log.Fatal("empty env GUSHER-CONN-TEST_SUBSCRIBE_MESSAGE")
	}
	sub_resp := os.Getenv("GUSHER-CONN-TEST_SUBSCRIBE_RESPONSE")
	if sub_resp == "" {
		log.Fatal("empty env GUSHER-CONN-TEST_SUBSCRIBE_RESPONSE")
	}
	push_msg := os.Getenv("GUSHER-CONN-TEST_PUSH_MESSAGE")
	if push_msg == "" {
		log.Fatal("empty env GUSHER-CONN-TEST_PUSH_MESSAGE")
	}
	connections := os.Getenv("GUSHER-CONN-TEST_CONNECTIONS")
	if connections == "" {
		log.Fatal("empty env GUSHER-CONN-TEST_CONNECTIONS")
	}
	conn_total, err := strconv.Atoi(connections)
	if err != nil {
		log.Fatal(err)
	}
	/*auth*/

	wsAuthurl, err := url.Parse(ws_auth_api)
	if err != nil {
		log.Fatal(err)
	}
	work := requestwork.New(5)
	loginUrl := url.Values{}

	loginUrl.Add("jwt", jwt)
	tokenChan := make(chan string, conn_total)
	tokenGroup := sync.WaitGroup{}
	for i := 0; i < conn_total; i++ {
		tokenGroup.Add(1)
		go func() {
			defer tokenGroup.Done()
			req, err := http.NewRequest("POST", wsAuthurl.String(), bytes.NewBufferString(loginUrl.Encode()))
			if err != nil {
				log.Fatal(err)
			}
			req.Header.Add("Content-Type", "application/x-www-form-urlencoded")
			req.Header.Add("Content-Length", strconv.Itoa(len(loginUrl.Encode())))
			ctx, _ := context.WithTimeout(context.Background(), 30*time.Second)
			err = work.Execute(ctx, req, func(resp *http.Response, e error) (err error) {
				if e != nil {
					return e
				}
				defer resp.Body.Close()
				b, err := ioutil.ReadAll(resp.Body)
				if err != nil {
					return err
				}
				v, err := jsonparser.GetString(b, "token")
				if err != nil {
					return err
				}
				tokenChan <- v
				return
			})
		}()

	}
	tokenGroup.Wait()
	close(tokenChan)
	/**/
	wsurlChan := make(chan *url.URL, conn_total)
	wsurlGroup := sync.WaitGroup{}
	for v := range tokenChan {
		wsurlGroup.Add(1)
		go func() {
			defer wsurlGroup.Done()
			wsurl, err := url.Parse(ws_api + "?token=" + v)
			if err != nil {
				return
			}
			wsurlChan <- wsurl

		}()
	}
	wsurlGroup.Wait()
	close(wsurlChan)
	pushurl, err := url.Parse(push_api)
	if err != nil {
		log.Fatal(err)
	}
	wsHeaders := http.Header{
		"Origin":                   {"*"},
		"Sec-WebSocket-Extensions": {"permessage-deflate; client_max_window_bits, x-webkit-deflate-frame"},
	}
	connChan := make(chan *websocket.Conn, conn_total)

	connGroup := sync.WaitGroup{}
	log.Infof("%v connect start!", conn_total)
	for wsurl := range wsurlChan {
		connGroup.Add(1)
		go func() {
			defer connGroup.Done()
			rawConn, err := net.Dial("tcp", wsurl.Host)
			if err != nil {
				return
			}

			conn, _, err := websocket.NewClient(rawConn, wsurl, wsHeaders, 8192, 8192)
			if err != nil {
				rawConn.Close()
				return
			}
			connChan <- conn
			return
		}()
	}
	connGroup.Wait()
	close(connChan)
	connTotal := len(connChan)
	var counter uint64
	for conn := range connChan {
		wg.Add(1)
		listen_wg.Add(1)
		go func(conn *websocket.Conn) {
			subStatus := false
			for {
				_, d, err := conn.ReadMessage()
				if err != nil {
					if c.Bool("debug") {
						log.Error(err)
					}
					if !subStatus {
						listen_wg.Done()
					}
					wg.Done()
					atomic.AddUint64(&counter, 1)
					return
				}

				if string(d) == sub_resp {
					subStatus = true
					listen_wg.Done()
				}
				if c.Bool("debug") {
					log.Println("slave repsonse message", string(d))
				}
				data, _ := jsonparser.GetString(d, "data")
				if data == push_msg {
					wg.Done()
					return
				}
			}
		}(conn)
		err = conn.WriteMessage(websocket.TextMessage, []byte(sub_msg))
		if err != nil {
			conn.Close()
			continue
		}
	}

	listen_wg.Wait()
	log.Infof("%v connect finish", connTotal)
	//push start
	v := url.Values{}

	v.Add("data", push_msg)
	req, err := http.NewRequest("POST", pushurl.String(), bytes.NewBufferString(v.Encode()))
	if err != nil {
		log.Fatal(err)
	}
	req.Header.Add("Content-Type", "application/x-www-form-urlencoded")
	req.Header.Add("Content-Length", strconv.Itoa(len(v.Encode())))

	if err != nil {
		log.Fatal(err)
	}
	var pushStart time.Time
	ctx, _ := context.WithTimeout(context.Background(), 30*time.Second)
	err = work.Execute(ctx, req, func(resp *http.Response, e error) (err error) {
		if e != nil {
			return
		}
		defer resp.Body.Close()
		b, err := ioutil.ReadAll(resp.Body)
		if err != nil {
			return err
		}
		if c.Bool("debug") {
			log.Println("master response", string(b))
		}
		pushStart = time.Now()
		return
	})
	log.Println("Waiting...")
	wg.Wait()
	t := time.Now().Sub(pushStart)
	if connTotal == 0 {
		log.Error("0 client connect, please check slave server!")
	} else if connTotal == int(counter) {
		log.Error("no client read message, please check master server!")
	} else {

		log.Infof("%v client connect, %v error read , receive msg time:%s", connTotal, counter, t)
	}

	return
}

func main() {
	cli.AppHelpTemplate += "\nWEBSITE:\n\t\thttps://github.com/syhlion/gusher.cluster/tree/master/test/conn-test\n\n"
	gusher := cli.NewApp()
	gusher.Usage = "simple connection test for gusher.cluster"
	gusher.Name = name
	gusher.Author = "Scott (syhlion)"
	gusher.Version = version
	gusher.Compiled = time.Now()
	gusher.Commands = []cli.Command{
		cmdStart,
	}
	gusher.Run(os.Args)
}
<commit_msg>修正 變數覆蓋問題<commit_after>package main

import (
	"bytes"
	"context"
	"io/ioutil"
	"net"
	"net/http"
	"net/url"
	"os"
	"strconv"
	"sync"
	"sync/atomic"
	"time"

	log "github.com/Sirupsen/logrus"
	"github.com/buger/jsonparser"
	"github.com/gorilla/websocket"
	"github.com/joho/godotenv"
	"github.com/syhlion/requestwork.v2"
	"github.com/urfave/cli"
)

var (
	name     string
	version  string
	cmdStart = cli.Command{
		Name:    "run",
		Aliases: []string{"r"},
		Usage:   "test connect to websocket ",
		Action:  start,
		Flags: []cli.Flag{
			cli.StringFlag{
				Name:  "env-file,e",
				Usage: "import env file",
			},
			cli.BoolFlag{
				Name:  "debug,d",
				Usage: "open debug mode",
			},
		},
	}
	wg        sync.WaitGroup
	listen_wg sync.WaitGroup
)

func start(c *cli.Context) {
	if c.String("env-file") != "" {
		envfile := c.String("env-file")
		//flag.Parse()
		err := godotenv.Load(envfile)
		if err != nil {
			log.Fatal(err)
		}
	}
	ws_api := os.Getenv("GUSHER-CONN-TEST_WS_API")
	if ws_api == "" {
		log.Fatal("empty env GUSHER-CONN-TEST_WS_API")
	}
	ws_auth_api := os.Getenv("GUSHER-CONN-TEST_WS_AUTH_API")
	if ws_auth_api == "" {
		log.Fatal("empty env GUSHER-CONN-TEST_WS_AUTH_API")
	}
	push_api := os.Getenv("GUSHER-CONN-TEST_PUSH_API")
	if push_api == "" {
		log.Fatal("empty env GUSHER-CONN-TEST_PUSH_API")
	}
	jwt := os.Getenv("GUSHER-CONN-TEST_JWT")
	if jwt == "" {
		log.Fatal("empty env GUSHER-CONN-TEST_JWT")
	}
	sub_msg := os.Getenv("GUSHER-CONN-TEST_SUBSCRIBE_MESSAGE")
	if sub_msg == "" {
		log.Fatal("empty env GUSHER-CONN-TEST_SUBSCRIBE_MESSAGE")
	}
	sub_resp := os.Getenv("GUSHER-CONN-TEST_SUBSCRIBE_RESPONSE")
	if sub_resp == "" {
		log.Fatal("empty env GUSHER-CONN-TEST_SUBSCRIBE_RESPONSE")
	}
	push_msg := os.Getenv("GUSHER-CONN-TEST_PUSH_MESSAGE")
	if push_msg == "" {
		log.Fatal("empty env GUSHER-CONN-TEST_PUSH_MESSAGE")
	}
	connections := os.Getenv("GUSHER-CONN-TEST_CONNECTIONS")
	if connections == "" {
		log.Fatal("empty env GUSHER-CONN-TEST_CONNECTIONS")
	}
	conn_total, err := strconv.Atoi(connections)
	if err != nil {
		log.Fatal(err)
	}
	/*auth*/

	wsAuthurl, err := url.Parse(ws_auth_api)
	if err != nil {
		log.Fatal(err)
	}
	work := requestwork.New(5)
	loginUrl := url.Values{}

	loginUrl.Add("jwt", jwt)
	tokenChan := make(chan string, conn_total)
	tokenGroup := sync.WaitGroup{}
	for i := 0; i < conn_total; i++ {
		tokenGroup.Add(1)
		go func() {
			defer tokenGroup.Done()
			req, err := http.NewRequest("POST", wsAuthurl.String(), bytes.NewBufferString(loginUrl.Encode()))
			if err != nil {
				log.Fatal(err)
			}
			req.Header.Add("Content-Type", "application/x-www-form-urlencoded")
			req.Header.Add("Content-Length", strconv.Itoa(len(loginUrl.Encode())))
			ctx, _ := context.WithTimeout(context.Background(), 30*time.Second)
			err = work.Execute(ctx, req, func(resp *http.Response, e error) (err error) {
				if e != nil {
					return e
				}
				defer resp.Body.Close()
				b, err := ioutil.ReadAll(resp.Body)
				if err != nil {
					return err
				}
				v, err := jsonparser.GetString(b, "token")
				if err != nil {
					return err
				}
				tokenChan <- v
				return
			})
		}()

	}
	tokenGroup.Wait()
	close(tokenChan)
	/**/
	wsurlChan := make(chan *url.URL, conn_total)
	wsurlGroup := sync.WaitGroup{}
	for v := range tokenChan {
		wsurlGroup.Add(1)
		go func(v string) {
			defer wsurlGroup.Done()
			wsurl, err := url.Parse(ws_api + "?token=" + v)
			if err != nil {
				return
			}
			wsurlChan <- wsurl

		}(v)
	}
	wsurlGroup.Wait()
	close(wsurlChan)
	pushurl, err := url.Parse(push_api)
	if err != nil {
		log.Fatal(err)
	}
	wsHeaders := http.Header{
		"Origin":                   {"*"},
		"Sec-WebSocket-Extensions": {"permessage-deflate; client_max_window_bits, x-webkit-deflate-frame"},
	}
	connChan := make(chan *websocket.Conn, conn_total)

	connGroup := sync.WaitGroup{}
	log.Infof("%v connect start!", conn_total)
	for wsurl := range wsurlChan {
		connGroup.Add(1)
		go func(wsurl *url.URL) {
			defer connGroup.Done()
			rawConn, err := net.Dial("tcp", wsurl.Host)
			if err != nil {
				return
			}

			conn, _, err := websocket.NewClient(rawConn, wsurl, wsHeaders, 8192, 8192)
			if err != nil {
				rawConn.Close()
				return
			}
			connChan <- conn
			return
		}(wsurl)
	}
	connGroup.Wait()
	close(connChan)
	connTotal := len(connChan)
	var counter uint64
	for conn := range connChan {
		wg.Add(1)
		listen_wg.Add(1)
		go func(conn *websocket.Conn) {
			subStatus := false
			for {
				_, d, err := conn.ReadMessage()
				if err != nil {
					if c.Bool("debug") {
						log.Error(err)
					}
					if !subStatus {
						listen_wg.Done()
					}
					wg.Done()
					atomic.AddUint64(&counter, 1)
					return
				}

				if string(d) == sub_resp {
					subStatus = true
					listen_wg.Done()
				}
				if c.Bool("debug") {
					log.Println("slave repsonse message", string(d))
				}
				data, _ := jsonparser.GetString(d, "data")
				if data == push_msg {
					wg.Done()
					return
				}
			}
		}(conn)
		err = conn.WriteMessage(websocket.TextMessage, []byte(sub_msg))
		if err != nil {
			conn.Close()
			continue
		}
	}

	listen_wg.Wait()
	log.Infof("%v connect finish", connTotal)
	//push start
	v := url.Values{}

	v.Add("data", push_msg)
	req, err := http.NewRequest("POST", pushurl.String(), bytes.NewBufferString(v.Encode()))
	if err != nil {
		log.Fatal(err)
	}
	req.Header.Add("Content-Type", "application/x-www-form-urlencoded")
	req.Header.Add("Content-Length", strconv.Itoa(len(v.Encode())))

	if err != nil {
		log.Fatal(err)
	}
	var pushStart time.Time
	ctx, _ := context.WithTimeout(context.Background(), 30*time.Second)
	err = work.Execute(ctx, req, func(resp *http.Response, e error) (err error) {
		if e != nil {
			return
		}
		defer resp.Body.Close()
		b, err := ioutil.ReadAll(resp.Body)
		if err != nil {
			return err
		}
		if c.Bool("debug") {
			log.Println("master response", string(b))
		}
		pushStart = time.Now()
		return
	})
	log.Println("Waiting...")
	wg.Wait()
	t := time.Now().Sub(pushStart)
	if connTotal == 0 {
		log.Error("0 client connect, please check slave server!")
	} else if connTotal == int(counter) {
		log.Error("no client read message, please check master server!")
	} else {

		log.Infof("%v client connect, %v error read , receive msg time:%s", connTotal, counter, t)
	}

	return
}

func main() {
	cli.AppHelpTemplate += "\nWEBSITE:\n\t\thttps://github.com/syhlion/gusher.cluster/tree/master/test/conn-test\n\n"
	gusher := cli.NewApp()
	gusher.Usage = "simple connection test for gusher.cluster"
	gusher.Name = name
	gusher.Author = "Scott (syhlion)"
	gusher.Version = version
	gusher.Compiled = time.Now()
	gusher.Commands = []cli.Command{
		cmdStart,
	}
	gusher.Run(os.Args)
}
<|endoftext|>


<commit_before>package db

import (
	"errors"
)

func InsertNewConnection(username string, connection string) (error) {
	db := GetDatabaseInstance()
	_, err := db.Exec("INSERT INTO connections VALUES (?, ?, 'pending')", username, connection)
	return err
}

func UpdateConnectionStatus(username string, connection string, status string) (error) {
	db := GetDatabaseInstance()
	_, err := db.Exec("UPDATE connections SET status = ? WHERE username = ? AND connection = ?",
												status, username, connection)
	return err
}

func DeleteConnection(username string, connection string) (int64, error) {
	db := GetDatabaseInstance()

  var affect int64
  qry, err := db.Prepare("DELETE FROM connections WHERE username = ? AND connection = ?")
  if err != nil {
    return affect, errors.New("Prepare Update Error!")
  }

  res, err := qry.Exec(username, connection)
  if err != nil {
    return affect, errors.New("Update Query Error!")
  }

  affect, err = res.RowsAffected()
  if err != nil {
    return affect, errors.New("Internal Update Error!")
  }

  return affect, nil
}

func GetConnectionsMap(user string) (map[string]bool, error) {
  return GetQueryResultsMap("SELECT connection FROM connections WHERE username = ?", user)
}
<commit_msg>added code in DeleteConnection to delete from either user in connection<commit_after>package db

import (
	"errors"
)

func InsertNewConnection(username string, connection string) (error) {
	db := GetDatabaseInstance()
	_, err := db.Exec("INSERT INTO connections VALUES (?, ?, 'pending')", username, connection)
	return err
}

func UpdateConnectionStatus(username string, connection string, status string) (error) {
	db := GetDatabaseInstance()
	_, err := db.Exec("UPDATE connections SET status = ? WHERE username = ? AND connection = ?",
												status, username, connection)
	return err
}

func DeleteConnection(username string, connection string) (int64, error) {
	db := GetDatabaseInstance()

  var affect int64
  primaryUser := username
  connectUser := connection

  userMap, err := GetConnectionsMap(connection)
  if err != nil {
    return affect, errors.New("Connection Search Error!")
  }
  if userMap[username] {
    primaryUser = connection
    connectUser = username
  }

  qry, err := db.Prepare("DELETE FROM connections WHERE username = ? AND connection = ?")
  if err != nil {
    return affect, errors.New("Prepare Update Error!")
  }

  res, err := qry.Exec(primaryUser, connectUser)
  if err != nil {
    return affect, errors.New("Update Query Error!")
  }

  affect, err = res.RowsAffected()
  if err != nil {
    return affect, errors.New("Internal Update Error!")
  }

  return affect, nil
}

func GetConnectionsMap(user string) (map[string]bool, error) {
  return GetQueryResultsMap("SELECT connection FROM connections WHERE username = ?", user)
}
<|endoftext|>


<commit_before>package rpc

import (
	"code.google.com/p/go.net/websocket"
	"koding/newkite/dnode"
	"time"
)

const origin = "http://localhost"
const redialDurationStart = 1 * time.Second
const redialDurationMax = 60 * time.Second

// Dial is a helper for creating a Client for just calling methods on the server.
// Do not use it if you want to handle methods on client side. Instead create a
// new Client, register your methods on Client.Dnode then call Client.Dial().
func Dial(url string, reconnect bool) (*Client, error) {
	c := NewClient()
	c.Reconnect = reconnect

	err := c.Dial(url)
	if err != nil {
		return nil, err
	}

	return c, nil
}

// Client is a dnode RPC client.
type Client struct {
	// Websocket connection
	Conn *websocket.Conn

	// Dnode message processor.
	Dnode *dnode.Dnode

	// Implements dnode.Transport interface
	tr *wsTransport

	// A space for saving/reading extra properties about this client.
	Properties map[string]interface{}

	// Dialled URL, used to re-connect again.
	url string

	// Should we reconnect if disconnected?
	Reconnect bool

	// Time to wait before redial connection.
	redialDuration time.Duration

	onConnectHandlers    []func()
	onDisconnectHandlers []func()
}

// NewClient returns a pointer to new Client.
// You need to call Dial() before interacting with the Server.
func NewClient() *Client {
	p := make(map[string]interface{})
	tr := &wsTransport{properties: p}
	c := &Client{
		Properties:     p,
		tr:             tr,
		Dnode:          dnode.New(tr),
		redialDuration: redialDurationStart,
	}
	tr.client = c
	return c
}

// Dial connects to the dnode server on "url" and starts a goroutine
// that processes incoming messages.
//
// Do not forget to register your handlers on Client.Dnode
// before calling Dial() to prevent race conditions.
func (c *Client) Dial(url string) error {
	c.url = url
	err := c.dial()
	if err != nil {
		return err
	}

	go c.run()
	return nil
}

// dial makes a single Dial() and run onConnectHandlers if connects.
func (c *Client) dial() error {
	ws, err := websocket.Dial(c.url, "", origin)
	if err != nil {
		return err
	}

	// We are connected
	c.Conn = ws
	c.tr.conn = ws

	// Reset the wait time.
	c.redialDuration = redialDurationStart

	// Must be run in a goroutine because a handler may wait a response from server.
	go c.connected()

	return nil
}

// DialForever connects to the server in background.
// If the connection drops, it reconnects again.
func (c *Client) DialForever(url string) {
	c.url = url
	go c.dialForever()
}

func (c *Client) dialForever() {
	for c.dial() != nil {
		if !c.Reconnect {
			return
		}

		c.sleep()
	}
	go c.run()
}

// run consumes incoming dnode messages. Reconnects if necessary.
func (c *Client) run() (err error) {
	for {
	running:
		err = c.Dnode.Run()
		c.disconnected()
	dialAgain:
		if !c.Reconnect {
			break
		}

		err = c.dial()
		if err != nil {
			c.sleep()
			goto dialAgain
		}

		goto running
	}

	return err
}

// sleep is used to wait for a while between dial retries.
// Each time it is called the redialDuration is incremented.
func (c *Client) sleep() {
	time.Sleep(c.redialDuration)
	c.redialDuration *= 2
	if c.redialDuration > redialDurationMax {
		c.redialDuration = redialDurationMax
	}
}

// Close closes the underlying websocket connection.
func (c *Client) Close() {
	c.Conn.Close()
}

// Call calls a method with args on the dnode server.
func (c *Client) Call(method string, args ...interface{}) (map[string]dnode.Path, error) {
	return c.Dnode.Call(method, args...)
}

// OnConnect registers a function to run on client connect.
func (c *Client) OnConnect(handler func()) {
	c.onConnectHandlers = append(c.onConnectHandlers, handler)
}

// OnDisconnect registers a function to run on client disconnect.
func (c *Client) OnDisconnect(handler func()) {
	c.onDisconnectHandlers = append(c.onDisconnectHandlers, handler)
}

// connected runs the registered connect handlers.
func (c *Client) connected() {
	for _, handler := range c.onConnectHandlers {
		go handler()
	}
}

// disconnected runs the registered disconnect handlers.
func (c *Client) disconnected() {
	for _, handler := range c.onDisconnectHandlers {
		go handler()
	}
}

// ------------
// Transport
// ------------

// wsTransport implements dnode.Transport interface.
type wsTransport struct {
	client     *Client
	conn       *websocket.Conn
	properties map[string]interface{}
}

func (t *wsTransport) Send(msg []byte) error {
	println("Sending...", string(msg))
	return websocket.Message.Send(t.conn, string(msg))
}

func (t *wsTransport) Receive() ([]byte, error) {
	println("Receiving...")
	var msg []byte
	err := websocket.Message.Receive(t.conn, &msg)
	println("Received:", string(msg))
	return msg, err
}

// RemoteAddr returns the host:port as string if server connection.
func (t *wsTransport) RemoteAddr() string {
	if t.conn.IsServerConn() {
		return t.conn.Request().RemoteAddr
	}
	return ""
}

func (t *wsTransport) Properties() map[string]interface{} {
	return t.properties
}

func (t *wsTransport) Client() interface{} {
	return t.client
}
<commit_msg>kite: style fixes<commit_after>package rpc

import (
	"code.google.com/p/go.net/websocket"
	"koding/newkite/dnode"
	"time"
)

const origin = "http://localhost"
const redialDurationStart = 1 * time.Second
const redialDurationMax = 60 * time.Second

// Dial is a helper for creating a Client for just calling methods on the server.
// Do not use it if you want to handle methods on client side. Instead create a
// new Client, register your methods on Client.Dnode then call Client.Dial().
func Dial(url string, reconnect bool) (*Client, error) {
	c := NewClient()
	c.Reconnect = reconnect

	err := c.Dial(url)
	if err != nil {
		return nil, err
	}

	return c, nil
}

// Client is a dnode RPC client.
type Client struct {
	// Websocket connection
	Conn *websocket.Conn

	// Dnode message processor.
	Dnode *dnode.Dnode

	// Implements dnode.Transport interface
	tr *wsTransport

	// A space for saving/reading extra properties about this client.
	Properties map[string]interface{}

	// Dialled URL, used to re-connect again.
	url string

	// Should we reconnect if disconnected?
	Reconnect bool

	// Time to wait before redial connection.
	redialDuration time.Duration

	onConnectHandlers    []func()
	onDisconnectHandlers []func()
}

// NewClient returns a pointer to new Client.
// You need to call Dial() before interacting with the Server.
func NewClient() *Client {
	p := make(map[string]interface{})

	tr := &wsTransport{properties: p}

	c := &Client{
		Properties:     p,
		tr:             tr,
		Dnode:          dnode.New(tr),
		redialDuration: redialDurationStart,
	}

	tr.client = c

	return c
}

// Dial connects to the dnode server on "url" and starts a goroutine
// that processes incoming messages.
//
// Do not forget to register your handlers on Client.Dnode
// before calling Dial() to prevent race conditions.
func (c *Client) Dial(url string) error {
	c.url = url
	err := c.dial()
	if err != nil {
		return err
	}

	go c.run()
	return nil
}

// dial makes a single Dial() and run onConnectHandlers if connects.
func (c *Client) dial() error {
	ws, err := websocket.Dial(c.url, "", origin)
	if err != nil {
		return err
	}

	// We are connected
	c.Conn = ws
	c.tr.conn = ws

	// Reset the wait time.
	c.redialDuration = redialDurationStart

	// Must be run in a goroutine because a handler may wait a response from server.
	go c.callOnConnectHandlers()

	return nil
}

// DialForever connects to the server in background.
// If the connection drops, it reconnects again.
func (c *Client) DialForever(url string) {
	c.url = url
	go c.dialForever()
}

func (c *Client) dialForever() {
	for c.dial() != nil {
		if !c.Reconnect {
			return
		}

		c.sleep()
	}
	go c.run()
}

// run consumes incoming dnode messages. Reconnects if necessary.
func (c *Client) run() (err error) {
	for {
	running:
		err = c.Dnode.Run()
		c.callOnDisconnectHandlers()
	dialAgain:
		if !c.Reconnect {
			break
		}

		err = c.dial()
		if err != nil {
			c.sleep()
			goto dialAgain
		}

		goto running
	}

	return err
}

// sleep is used to wait for a while between dial retries.
// Each time it is called the redialDuration is incremented.
func (c *Client) sleep() {
	time.Sleep(c.redialDuration)

	c.redialDuration *= 2
	if c.redialDuration > redialDurationMax {
		c.redialDuration = redialDurationMax
	}
}

// Close closes the underlying websocket connection.
func (c *Client) Close() {
	c.Conn.Close()
}

// Call calls a method with args on the dnode server.
func (c *Client) Call(method string, args ...interface{}) (map[string]dnode.Path, error) {
	return c.Dnode.Call(method, args...)
}

// OnConnect registers a function to run on client connect.
func (c *Client) OnConnect(handler func()) {
	c.onConnectHandlers = append(c.onConnectHandlers, handler)
}

// OnDisconnect registers a function to run on client disconnect.
func (c *Client) OnDisconnect(handler func()) {
	c.onDisconnectHandlers = append(c.onDisconnectHandlers, handler)
}

// callOnConnectHandlers runs the registered connect handlers.
func (c *Client) callOnConnectHandlers() {
	for _, handler := range c.onConnectHandlers {
		go handler()
	}
}

// callOnDisconnectHandlers runs the registered disconnect handlers.
func (c *Client) callOnDisconnectHandlers() {
	for _, handler := range c.onDisconnectHandlers {
		go handler()
	}
}

// ------------
// Transport
// ------------

// wsTransport implements dnode.Transport interface.
type wsTransport struct {
	client     *Client
	conn       *websocket.Conn
	properties map[string]interface{}
}

func (t *wsTransport) Send(msg []byte) error {
	println("Sending...", string(msg))
	return websocket.Message.Send(t.conn, string(msg))
}

func (t *wsTransport) Receive() ([]byte, error) {
	println("Receiving...")
	var msg []byte
	err := websocket.Message.Receive(t.conn, &msg)
	println("Received:", string(msg))
	return msg, err
}

// RemoteAddr returns the host:port as string if server connection.
func (t *wsTransport) RemoteAddr() string {
	if t.conn.IsServerConn() {
		return t.conn.Request().RemoteAddr
	}
	return ""
}

func (t *wsTransport) Properties() map[string]interface{} {
	return t.properties
}

func (t *wsTransport) Client() interface{} {
	return t.client
}
<|endoftext|>


<commit_before>package scraper

import "testing"

func TestScrape(t *testing.T) {
	getItem = func(url string) {
		defer wg.Done()

		ch <- Item{
			"FooTitle",
			"FooSize",
			"10.00",
			"FooDescription",
		}
	}

	urls := []string{
		"http://foo.com/",
		"http://bar.com/",
		"http://baz.com/",
	}

	result := Scrape(urls)
	response := result.Total
	expected := "30.00"

	if response != expected {
		t.Errorf("The response:\n '%s'\ndidn't match the expectation:\n '%s'", response, expected)
	}
}
<commit_msg>Rename Test to be more descriptive<commit_after>package scraper

import "testing"

func TestScrapeResultsTotal(t *testing.T) {
	getItem = func(url string) {
		defer wg.Done()

		ch <- Item{
			"FooTitle",
			"FooSize",
			"10.00",
			"FooDescription",
		}
	}

	urls := []string{
		"http://foo.com/",
		"http://bar.com/",
		"http://baz.com/",
	}

	result := Scrape(urls)
	response := result.Total
	expected := "30.00"

	if response != expected {
		t.Errorf("The response:\n '%s'\ndidn't match the expectation:\n '%s'", response, expected)
	}
}
<|endoftext|>


<commit_before>package main

import (
	"encoding/json"
	"fmt"
	"github.com/gorilla/mux"
	"io"
	"koding/kontrol/kontroldaemon/workerconfig"
	"labix.org/v2/mgo/bson"
	"log"
	"net/http"
	"net/url"
	"sort"
	"strconv"
	"strings"
	"time"
)

type Worker struct {
	Name      string    `json:"name"`
	Uuid      string    `json:"uuid"`
	Hostname  string    `json:"hostname"`
	Version   int       `json:"version"`
	Timestamp time.Time `json:"timestamp"`
	Pid       int       `json:"pid"`
	State     string    `json:"state"`
	Uptime    int       `json:"uptime"`
	Port      int       `json:"port"`
}

type Workers []Worker

var StatusCode = map[workerconfig.WorkerStatus]string{
	workerconfig.Running:    "running",
	workerconfig.Pending:    "waiting",
	workerconfig.Waiting:    "waiting",
	workerconfig.Stopped:    "stopped",
	workerconfig.Notstarted: "stopped",
	workerconfig.Killed:     "dead",
	workerconfig.Dead:       "dead",
}

func GetWorkers(writer http.ResponseWriter, req *http.Request) {
	log.Println("GET /workers")
	queries, _ := url.ParseQuery(req.URL.RawQuery)

	var latestVersion bool
	query := bson.M{}
	for key, value := range queries {
		switch key {
		case "version", "pid":
			if value[0] == "latest" {
				latestVersion = true
				break
			}
			v, _ := strconv.Atoi(value[0])
			query[key] = v
		case "state":
			for status, state := range StatusCode {
				if value[0] == state {
					query["status"] = status
				}
			}
		default:
			if key == "name" {
				name := value[0]
				if counts := strings.Count(value[0], "-"); counts > 0 {
					s := strings.Split(value[0], "-")
					name = s[0]
				}
				query[key] = bson.RegEx{Pattern: "^" + name, Options: "i"}
			} else {
				query[key] = value[0]
			}
		}
	}

	matchedWorkers := queryResult(query, latestVersion)
	data, err := json.MarshalIndent(matchedWorkers, "", "  ")
	if err != nil {
		io.WriteString(writer, fmt.Sprintf("{\"err\":\"%s\"}\n", err))
		return
	}
	writer.Write(data)

}

func GetWorker(writer http.ResponseWriter, req *http.Request) {
	vars := mux.Vars(req)
	uuid := vars["uuid"]
	log.Printf("GET /workers/%s\n", uuid)

	query := bson.M{"uuid": uuid}
	matchedWorkers := queryResult(query, false)
	data, err := json.MarshalIndent(matchedWorkers, "", "  ")
	if err != nil {
		io.WriteString(writer, fmt.Sprintf("{\"err\":\"%s\"}\n", err))
		return
	}
	writer.Write(data)
}

func UpdateWorker(writer http.ResponseWriter, req *http.Request) {
	vars := mux.Vars(req)
	uuid, action := vars["uuid"], vars["action"]
	log.Printf("%s /workers/%s\n", strings.ToUpper(action), uuid)

	buildSendCmd(action, "", uuid)
	resp := fmt.Sprintf("worker: '%s' is updated in db", uuid)
	io.WriteString(writer, resp)
}

func DeleteWorker(writer http.ResponseWriter, req *http.Request) {
	vars := mux.Vars(req)
	uuid := vars["uuid"]
	log.Printf("DELETE /workers/%s\n", uuid)

	buildSendCmd("delete", "", uuid)
	resp := fmt.Sprintf("worker: '%s' is deleted from db", uuid)
	io.WriteString(writer, resp)
}

func queryResult(query bson.M, latestVersion bool) Workers {
	workers := make(Workers, 0)
	worker := workerconfig.MsgWorker{}

	iter := kontrolConfig.Collection.Find(query).Iter()
	for iter.Next(&worker) {
		apiWorker := &Worker{
			worker.Name,
			worker.Uuid,
			worker.Hostname,
			worker.Version,
			worker.Timestamp,
			worker.Pid,
			StatusCode[worker.Status],
			worker.Monitor.Uptime,
			worker.Port,
		}

		workers = append(workers, *apiWorker)
	}

	// finding the largest number of a field in mongo is kinda problematic.
	// therefore we are doing it on our side
	if latestVersion {
		versions := make([]int, len(workers))

		if len(workers) == 0 {
			return workers
		}

		for i, val := range workers {
			versions[i] = val.Version
		}

		sort.Ints(versions)
		maxVersion := versions[len(versions)-1] // get largest version number

		filteredWorkers := make(Workers, 0)
		for _, val := range workers {
			if maxVersion == val.Version {
				filteredWorkers = append(filteredWorkers, val)
			}
		}

		return filteredWorkers
	}

	return workers
}

func buildSendCmd(action, host, uuid string) {
	cmd := workerconfig.Request{Command: action, Hostname: host, Uuid: uuid}
	log.Println("Sending cmd to kontrold:", cmd)

	// Wrap message for dynamic unmarshaling at endpoint
	type Wrap struct{ Worker workerconfig.Request }

	data, err := json.Marshal(&Wrap{cmd})
	if err != nil {
		log.Println("Json marshall error", data)
	}

	amqpWrapper.Publish(data)
}
<commit_msg>kontrolapi: just a small note why use regex here<commit_after>package main

import (
	"encoding/json"
	"fmt"
	"github.com/gorilla/mux"
	"io"
	"koding/kontrol/kontroldaemon/workerconfig"
	"labix.org/v2/mgo/bson"
	"log"
	"net/http"
	"net/url"
	"sort"
	"strconv"
	"strings"
	"time"
)

type Worker struct {
	Name      string    `json:"name"`
	Uuid      string    `json:"uuid"`
	Hostname  string    `json:"hostname"`
	Version   int       `json:"version"`
	Timestamp time.Time `json:"timestamp"`
	Pid       int       `json:"pid"`
	State     string    `json:"state"`
	Uptime    int       `json:"uptime"`
	Port      int       `json:"port"`
}

type Workers []Worker

var StatusCode = map[workerconfig.WorkerStatus]string{
	workerconfig.Running:    "running",
	workerconfig.Pending:    "waiting",
	workerconfig.Waiting:    "waiting",
	workerconfig.Stopped:    "stopped",
	workerconfig.Notstarted: "stopped",
	workerconfig.Killed:     "dead",
	workerconfig.Dead:       "dead",
}

func GetWorkers(writer http.ResponseWriter, req *http.Request) {
	log.Println("GET /workers")
	queries, _ := url.ParseQuery(req.URL.RawQuery)

	var latestVersion bool
	query := bson.M{}
	for key, value := range queries {
		switch key {
		case "version", "pid":
			if value[0] == "latest" {
				latestVersion = true
				break
			}
			v, _ := strconv.Atoi(value[0])
			query[key] = v
		case "state":
			for status, state := range StatusCode {
				if value[0] == state {
					query["status"] = status
				}
			}
		default:
			if key == "name" {
				name := value[0]
				if counts := strings.Count(value[0], "-"); counts > 0 {
					s := strings.Split(value[0], "-")
					name = s[0]
				}
				// if searched for social-1, social-2, then return all workers
				// that begins with social
				query[key] = bson.RegEx{Pattern: "^" + name, Options: "i"}
			} else {
				query[key] = value[0]
			}
		}
	}

	matchedWorkers := queryResult(query, latestVersion)
	data, err := json.MarshalIndent(matchedWorkers, "", "  ")
	if err != nil {
		io.WriteString(writer, fmt.Sprintf("{\"err\":\"%s\"}\n", err))
		return
	}
	writer.Write(data)

}

func GetWorker(writer http.ResponseWriter, req *http.Request) {
	vars := mux.Vars(req)
	uuid := vars["uuid"]
	log.Printf("GET /workers/%s\n", uuid)

	query := bson.M{"uuid": uuid}
	matchedWorkers := queryResult(query, false)
	data, err := json.MarshalIndent(matchedWorkers, "", "  ")
	if err != nil {
		io.WriteString(writer, fmt.Sprintf("{\"err\":\"%s\"}\n", err))
		return
	}
	writer.Write(data)
}

func UpdateWorker(writer http.ResponseWriter, req *http.Request) {
	vars := mux.Vars(req)
	uuid, action := vars["uuid"], vars["action"]
	log.Printf("%s /workers/%s\n", strings.ToUpper(action), uuid)

	buildSendCmd(action, "", uuid)
	resp := fmt.Sprintf("worker: '%s' is updated in db", uuid)
	io.WriteString(writer, resp)
}

func DeleteWorker(writer http.ResponseWriter, req *http.Request) {
	vars := mux.Vars(req)
	uuid := vars["uuid"]
	log.Printf("DELETE /workers/%s\n", uuid)

	buildSendCmd("delete", "", uuid)
	resp := fmt.Sprintf("worker: '%s' is deleted from db", uuid)
	io.WriteString(writer, resp)
}

func queryResult(query bson.M, latestVersion bool) Workers {
	workers := make(Workers, 0)
	worker := workerconfig.MsgWorker{}

	iter := kontrolConfig.Collection.Find(query).Iter()
	for iter.Next(&worker) {
		apiWorker := &Worker{
			worker.Name,
			worker.Uuid,
			worker.Hostname,
			worker.Version,
			worker.Timestamp,
			worker.Pid,
			StatusCode[worker.Status],
			worker.Monitor.Uptime,
			worker.Port,
		}

		workers = append(workers, *apiWorker)
	}

	// finding the largest number of a field in mongo is kinda problematic.
	// therefore we are doing it on our side
	if latestVersion {
		versions := make([]int, len(workers))

		if len(workers) == 0 {
			return workers
		}

		for i, val := range workers {
			versions[i] = val.Version
		}

		sort.Ints(versions)
		maxVersion := versions[len(versions)-1] // get largest version number

		filteredWorkers := make(Workers, 0)
		for _, val := range workers {
			if maxVersion == val.Version {
				filteredWorkers = append(filteredWorkers, val)
			}
		}

		return filteredWorkers
	}

	return workers
}

func buildSendCmd(action, host, uuid string) {
	cmd := workerconfig.Request{Command: action, Hostname: host, Uuid: uuid}
	log.Println("Sending cmd to kontrold:", cmd)

	// Wrap message for dynamic unmarshaling at endpoint
	type Wrap struct{ Worker workerconfig.Request }

	data, err := json.Marshal(&Wrap{cmd})
	if err != nil {
		log.Println("Json marshall error", data)
	}

	amqpWrapper.Publish(data)
}
<|endoftext|>


<commit_before>/*
Copyright 2017 The Kubernetes Authors.

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

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

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

package apiserver

import (
	"fmt"
	"net/http"
	"strings"
	"time"

	"k8s.io/apiextensions-apiserver/pkg/apis/apiextensions"
	"k8s.io/apiextensions-apiserver/pkg/apis/apiextensions/install"
	"k8s.io/apiextensions-apiserver/pkg/apis/apiextensions/v1"
	"k8s.io/apiextensions-apiserver/pkg/apis/apiextensions/v1beta1"
	_ "k8s.io/apiextensions-apiserver/pkg/client/clientset/clientset"
	"k8s.io/apiextensions-apiserver/pkg/client/clientset/internalclientset"
	_ "k8s.io/apiextensions-apiserver/pkg/client/informers/externalversions"
	internalinformers "k8s.io/apiextensions-apiserver/pkg/client/informers/internalversion"
	"k8s.io/apiextensions-apiserver/pkg/controller/apiapproval"
	"k8s.io/apiextensions-apiserver/pkg/controller/establish"
	"k8s.io/apiextensions-apiserver/pkg/controller/finalizer"
	"k8s.io/apiextensions-apiserver/pkg/controller/nonstructuralschema"
	openapicontroller "k8s.io/apiextensions-apiserver/pkg/controller/openapi"
	"k8s.io/apiextensions-apiserver/pkg/controller/status"
	apiextensionsfeatures "k8s.io/apiextensions-apiserver/pkg/features"
	"k8s.io/apiextensions-apiserver/pkg/registry/customresourcedefinition"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/apimachinery/pkg/labels"
	"k8s.io/apimachinery/pkg/runtime"
	"k8s.io/apimachinery/pkg/runtime/schema"
	"k8s.io/apimachinery/pkg/runtime/serializer"
	"k8s.io/apimachinery/pkg/util/sets"
	"k8s.io/apimachinery/pkg/util/wait"
	"k8s.io/apimachinery/pkg/version"
	"k8s.io/apiserver/pkg/endpoints/discovery"
	genericregistry "k8s.io/apiserver/pkg/registry/generic"
	"k8s.io/apiserver/pkg/registry/rest"
	genericapiserver "k8s.io/apiserver/pkg/server"
	serverstorage "k8s.io/apiserver/pkg/server/storage"
	utilfeature "k8s.io/apiserver/pkg/util/feature"
	"k8s.io/apiserver/pkg/util/webhook"
	"k8s.io/klog"
)

var (
	Scheme = runtime.NewScheme()
	Codecs = serializer.NewCodecFactory(Scheme)

	// if you modify this, make sure you update the crEncoder
	unversionedVersion = schema.GroupVersion{Group: "", Version: "v1"}
	unversionedTypes   = []runtime.Object{
		&metav1.Status{},
		&metav1.WatchEvent{},
		&metav1.APIVersions{},
		&metav1.APIGroupList{},
		&metav1.APIGroup{},
		&metav1.APIResourceList{},
	}
)

func init() {
	install.Install(Scheme)

	// we need to add the options to empty v1
	metav1.AddToGroupVersion(Scheme, schema.GroupVersion{Group: "", Version: "v1"})

	Scheme.AddUnversionedTypes(unversionedVersion, unversionedTypes...)
}

type ExtraConfig struct {
	CRDRESTOptionsGetter genericregistry.RESTOptionsGetter

	// MasterCount is used to detect whether cluster is HA, and if it is
	// the CRD Establishing will be hold by 5 seconds.
	MasterCount int

	// ServiceResolver is used in CR webhook converters to resolve webhook's service names
	ServiceResolver webhook.ServiceResolver
	// AuthResolverWrapper is used in CR webhook converters
	AuthResolverWrapper webhook.AuthenticationInfoResolverWrapper
}

type Config struct {
	GenericConfig *genericapiserver.RecommendedConfig
	ExtraConfig   ExtraConfig
}

type completedConfig struct {
	GenericConfig genericapiserver.CompletedConfig
	ExtraConfig   *ExtraConfig
}

type CompletedConfig struct {
	// Embed a private pointer that cannot be instantiated outside of this package.
	*completedConfig
}

type CustomResourceDefinitions struct {
	GenericAPIServer *genericapiserver.GenericAPIServer

	// provided for easier embedding
	Informers internalinformers.SharedInformerFactory
}

// Complete fills in any fields not set that are required to have valid data. It's mutating the receiver.
func (cfg *Config) Complete() CompletedConfig {
	c := completedConfig{
		cfg.GenericConfig.Complete(),
		&cfg.ExtraConfig,
	}

	c.GenericConfig.EnableDiscovery = false
	c.GenericConfig.Version = &version.Info{
		Major: "0",
		Minor: "1",
	}

	return CompletedConfig{&c}
}

// New returns a new instance of CustomResourceDefinitions from the given config.
func (c completedConfig) New(delegationTarget genericapiserver.DelegationTarget) (*CustomResourceDefinitions, error) {
	genericServer, err := c.GenericConfig.New("apiextensions-apiserver", delegationTarget)
	if err != nil {
		return nil, err
	}

	s := &CustomResourceDefinitions{
		GenericAPIServer: genericServer,
	}

	apiResourceConfig := c.GenericConfig.MergedResourceConfig
	apiGroupInfo := genericapiserver.NewDefaultAPIGroupInfo(apiextensions.GroupName, Scheme, metav1.ParameterCodec, Codecs)
	if apiResourceConfig.VersionEnabled(v1beta1.SchemeGroupVersion) {
		storage := map[string]rest.Storage{}
		// customresourcedefinitions
		customResourceDefintionStorage := customresourcedefinition.NewREST(Scheme, c.GenericConfig.RESTOptionsGetter)
		storage["customresourcedefinitions"] = customResourceDefintionStorage
		storage["customresourcedefinitions/status"] = customresourcedefinition.NewStatusREST(Scheme, customResourceDefintionStorage)

		apiGroupInfo.VersionedResourcesStorageMap[v1beta1.SchemeGroupVersion.Version] = storage
	}
	if apiResourceConfig.VersionEnabled(v1.SchemeGroupVersion) {
		storage := map[string]rest.Storage{}
		// customresourcedefinitions
		customResourceDefintionStorage := customresourcedefinition.NewREST(Scheme, c.GenericConfig.RESTOptionsGetter)
		storage["customresourcedefinitions"] = customResourceDefintionStorage
		storage["customresourcedefinitions/status"] = customresourcedefinition.NewStatusREST(Scheme, customResourceDefintionStorage)

		apiGroupInfo.VersionedResourcesStorageMap[v1.SchemeGroupVersion.Version] = storage
	}

	if err := s.GenericAPIServer.InstallAPIGroup(&apiGroupInfo); err != nil {
		return nil, err
	}

	crdClient, err := internalclientset.NewForConfig(s.GenericAPIServer.LoopbackClientConfig)
	if err != nil {
		// it's really bad that this is leaking here, but until we can fix the test (which I'm pretty sure isn't even testing what it wants to test),
		// we need to be able to move forward
		return nil, fmt.Errorf("failed to create clientset: %v", err)
	}
	s.Informers = internalinformers.NewSharedInformerFactory(crdClient, 5*time.Minute)

	delegateHandler := delegationTarget.UnprotectedHandler()
	if delegateHandler == nil {
		delegateHandler = http.NotFoundHandler()
	}

	versionDiscoveryHandler := &versionDiscoveryHandler{
		discovery: map[schema.GroupVersion]*discovery.APIVersionHandler{},
		delegate:  delegateHandler,
	}
	groupDiscoveryHandler := &groupDiscoveryHandler{
		discovery: map[string]*discovery.APIGroupHandler{},
		delegate:  delegateHandler,
	}
	establishingController := establish.NewEstablishingController(s.Informers.Apiextensions().InternalVersion().CustomResourceDefinitions(), crdClient.Apiextensions())
	crdHandler, err := NewCustomResourceDefinitionHandler(
		versionDiscoveryHandler,
		groupDiscoveryHandler,
		s.Informers.Apiextensions().InternalVersion().CustomResourceDefinitions(),
		delegateHandler,
		c.ExtraConfig.CRDRESTOptionsGetter,
		c.GenericConfig.AdmissionControl,
		establishingController,
		c.ExtraConfig.ServiceResolver,
		c.ExtraConfig.AuthResolverWrapper,
		c.ExtraConfig.MasterCount,
		s.GenericAPIServer.Authorizer,
		c.GenericConfig.RequestTimeout,
		time.Duration(c.GenericConfig.MinRequestTimeout)*time.Second,
		apiGroupInfo.StaticOpenAPISpec,
		c.GenericConfig.MaxRequestBodyBytes,
	)
	if err != nil {
		return nil, err
	}
	s.GenericAPIServer.Handler.NonGoRestfulMux.Handle("/apis", crdHandler)
	s.GenericAPIServer.Handler.NonGoRestfulMux.HandlePrefix("/apis/", crdHandler)

	crdController := NewDiscoveryController(s.Informers.Apiextensions().InternalVersion().CustomResourceDefinitions(), versionDiscoveryHandler, groupDiscoveryHandler)
	namingController := status.NewNamingConditionController(s.Informers.Apiextensions().InternalVersion().CustomResourceDefinitions(), crdClient.Apiextensions())
	nonStructuralSchemaController := nonstructuralschema.NewConditionController(s.Informers.Apiextensions().InternalVersion().CustomResourceDefinitions(), crdClient.Apiextensions())
	apiApprovalController := apiapproval.NewKubernetesAPIApprovalPolicyConformantConditionController(s.Informers.Apiextensions().InternalVersion().CustomResourceDefinitions(), crdClient.Apiextensions())
	finalizingController := finalizer.NewCRDFinalizer(
		s.Informers.Apiextensions().InternalVersion().CustomResourceDefinitions(),
		crdClient.Apiextensions(),
		crdHandler,
	)
	var openapiController *openapicontroller.Controller
	if utilfeature.DefaultFeatureGate.Enabled(apiextensionsfeatures.CustomResourcePublishOpenAPI) {
		openapiController = openapicontroller.NewController(s.Informers.Apiextensions().InternalVersion().CustomResourceDefinitions())
	}

	s.GenericAPIServer.AddPostStartHookOrDie("start-apiextensions-informers", func(context genericapiserver.PostStartHookContext) error {
		s.Informers.Start(context.StopCh)
		return nil
	})
	s.GenericAPIServer.AddPostStartHookOrDie("start-apiextensions-controllers", func(context genericapiserver.PostStartHookContext) error {
		// OpenAPIVersionedService and StaticOpenAPISpec are populated in generic apiserver PrepareRun().
		// Together they serve the /openapi/v2 endpoint on a generic apiserver. A generic apiserver may
		// choose to not enable OpenAPI by having null openAPIConfig, and thus OpenAPIVersionedService
		// and StaticOpenAPISpec are both null. In that case we don't run the CRD OpenAPI controller.
		if utilfeature.DefaultFeatureGate.Enabled(apiextensionsfeatures.CustomResourcePublishOpenAPI) && s.GenericAPIServer.OpenAPIVersionedService != nil && s.GenericAPIServer.StaticOpenAPISpec != nil {
			go openapiController.Run(s.GenericAPIServer.StaticOpenAPISpec, s.GenericAPIServer.OpenAPIVersionedService, context.StopCh)
		}

		go crdController.Run(context.StopCh)
		go namingController.Run(context.StopCh)
		go establishingController.Run(context.StopCh)
		go nonStructuralSchemaController.Run(5, context.StopCh)
		go apiApprovalController.Run(5, context.StopCh)
		go finalizingController.Run(5, context.StopCh)
		return nil
	})
	s.GenericAPIServer.AddPostStartHookOrDie("crd-discovery-available", func(context genericapiserver.PostStartHookContext) error {
		return wait.PollImmediateUntil(100*time.Millisecond, func() (bool, error) {
			// only check if we have a valid list for a given resourceversion
			if !s.Informers.Apiextensions().InternalVersion().CustomResourceDefinitions().Informer().HasSynced() {
				return false, nil
			}

			// The returned group and resource lists might be non-nil with partial results even in the
			// case of non-nil error.  If API aggregation fails, we still want our other discovery information because the CRDs
			// may all be present.
			_, serverGroupsAndResources, discoveryErr := crdClient.Discovery().ServerGroupsAndResources()
			if discoveryErr != nil {
				klog.V(2).Info(discoveryErr)
			}

			serverCRDs, err := s.Informers.Apiextensions().InternalVersion().CustomResourceDefinitions().Lister().List(labels.Everything())
			if err != nil {
				return false, err
			}
			crdGroupsAndResources := sets.NewString()
			for _, crd := range serverCRDs {
				// Skip not active CRD
				if !apiextensions.IsCRDConditionTrue(crd, apiextensions.Established) {
					continue
				}
				for _, version := range crd.Spec.Versions {
					// Skip versions that are not served
					if !version.Served {
						continue
					}
					crdGroupsAndResources.Insert(fmt.Sprintf("%s.%s.%s", crd.Spec.Names.Plural, version.Name, crd.Spec.Group))
				}
			}

			discoveryGroupsAndResources := sets.NewString()
			for _, resourceList := range serverGroupsAndResources {
				for _, apiResource := range resourceList.APIResources {
					group, version := splitGroupVersion(resourceList.GroupVersion)
					discoveryGroupsAndResources.Insert(fmt.Sprintf("%s.%s.%s", apiResource.Name, version, group))
				}
			}
			if !discoveryGroupsAndResources.HasAll(crdGroupsAndResources.List()...) {
				klog.Infof("waiting for CRD resources in discovery: %#v", crdGroupsAndResources.Difference(discoveryGroupsAndResources))
				return false, nil
			}
			return true, nil
		}, context.StopCh)
	})

	// we don't want to report healthy until we can handle all CRDs that have already been registered.  Waiting for the informer
	// to sync makes sure that the lister will be valid before we begin.  There may still be races for CRDs added after startup,
	// but we won't go healthy until we can handle the ones already present.
	s.GenericAPIServer.AddPostStartHookOrDie("crd-informer-synced", func(context genericapiserver.PostStartHookContext) error {
		return wait.PollImmediateUntil(100*time.Millisecond, func() (bool, error) {
			return s.Informers.Apiextensions().InternalVersion().CustomResourceDefinitions().Informer().HasSynced(), nil
		}, context.StopCh)
	})

	return s, nil
}

func splitGroupVersion(gv string) (group string, version string) {
	ss := strings.SplitN(gv, "/", 2)
	if len(ss) == 1 {
		version = ss[0]
	} else {
		group, version = ss[0], ss[1]
	}
	return
}

func DefaultAPIResourceConfigSource() *serverstorage.ResourceConfig {
	ret := serverstorage.NewResourceConfig()
	// NOTE: GroupVersions listed here will be enabled by default. Don't put alpha versions in the list.
	ret.EnableVersions(
		v1beta1.SchemeGroupVersion,
		v1.SchemeGroupVersion,
	)

	return ret
}
<commit_msg>UPSTREAM: <carry>: skip openshift-apiserver served CRDs from crd-discovery-available posthook<commit_after>/*
Copyright 2017 The Kubernetes Authors.

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

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

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

package apiserver

import (
	"fmt"
	"net/http"
	"strings"
	"time"

	"k8s.io/apiextensions-apiserver/pkg/apis/apiextensions"
	"k8s.io/apiextensions-apiserver/pkg/apis/apiextensions/install"
	"k8s.io/apiextensions-apiserver/pkg/apis/apiextensions/v1"
	"k8s.io/apiextensions-apiserver/pkg/apis/apiextensions/v1beta1"
	_ "k8s.io/apiextensions-apiserver/pkg/client/clientset/clientset"
	"k8s.io/apiextensions-apiserver/pkg/client/clientset/internalclientset"
	_ "k8s.io/apiextensions-apiserver/pkg/client/informers/externalversions"
	internalinformers "k8s.io/apiextensions-apiserver/pkg/client/informers/internalversion"
	"k8s.io/apiextensions-apiserver/pkg/controller/apiapproval"
	"k8s.io/apiextensions-apiserver/pkg/controller/establish"
	"k8s.io/apiextensions-apiserver/pkg/controller/finalizer"
	"k8s.io/apiextensions-apiserver/pkg/controller/nonstructuralschema"
	openapicontroller "k8s.io/apiextensions-apiserver/pkg/controller/openapi"
	"k8s.io/apiextensions-apiserver/pkg/controller/status"
	apiextensionsfeatures "k8s.io/apiextensions-apiserver/pkg/features"
	"k8s.io/apiextensions-apiserver/pkg/registry/customresourcedefinition"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/apimachinery/pkg/labels"
	"k8s.io/apimachinery/pkg/runtime"
	"k8s.io/apimachinery/pkg/runtime/schema"
	"k8s.io/apimachinery/pkg/runtime/serializer"
	"k8s.io/apimachinery/pkg/util/sets"
	"k8s.io/apimachinery/pkg/util/wait"
	"k8s.io/apimachinery/pkg/version"
	"k8s.io/apiserver/pkg/endpoints/discovery"
	genericregistry "k8s.io/apiserver/pkg/registry/generic"
	"k8s.io/apiserver/pkg/registry/rest"
	genericapiserver "k8s.io/apiserver/pkg/server"
	serverstorage "k8s.io/apiserver/pkg/server/storage"
	utilfeature "k8s.io/apiserver/pkg/util/feature"
	"k8s.io/apiserver/pkg/util/webhook"
	"k8s.io/klog"
)

var (
	Scheme = runtime.NewScheme()
	Codecs = serializer.NewCodecFactory(Scheme)

	// if you modify this, make sure you update the crEncoder
	unversionedVersion = schema.GroupVersion{Group: "", Version: "v1"}
	unversionedTypes   = []runtime.Object{
		&metav1.Status{},
		&metav1.WatchEvent{},
		&metav1.APIVersions{},
		&metav1.APIGroupList{},
		&metav1.APIGroup{},
		&metav1.APIResourceList{},
	}

	// a list of CRDs served by Openshift API Server, which we can't wait for
	// otherwise bootstrap fails, so we'll always ignore them.
	skipOpenshiftAPIServerCRDs = sets.NewString(
		"clusterresourcequotas.quota.openshift.io",
		"rolebindingrestrictions.authorization.openshift.io",
		"securitycontextconstraints.security.openshift.io",
	)
)

func init() {
	install.Install(Scheme)

	// we need to add the options to empty v1
	metav1.AddToGroupVersion(Scheme, schema.GroupVersion{Group: "", Version: "v1"})

	Scheme.AddUnversionedTypes(unversionedVersion, unversionedTypes...)
}

type ExtraConfig struct {
	CRDRESTOptionsGetter genericregistry.RESTOptionsGetter

	// MasterCount is used to detect whether cluster is HA, and if it is
	// the CRD Establishing will be hold by 5 seconds.
	MasterCount int

	// ServiceResolver is used in CR webhook converters to resolve webhook's service names
	ServiceResolver webhook.ServiceResolver
	// AuthResolverWrapper is used in CR webhook converters
	AuthResolverWrapper webhook.AuthenticationInfoResolverWrapper
}

type Config struct {
	GenericConfig *genericapiserver.RecommendedConfig
	ExtraConfig   ExtraConfig
}

type completedConfig struct {
	GenericConfig genericapiserver.CompletedConfig
	ExtraConfig   *ExtraConfig
}

type CompletedConfig struct {
	// Embed a private pointer that cannot be instantiated outside of this package.
	*completedConfig
}

type CustomResourceDefinitions struct {
	GenericAPIServer *genericapiserver.GenericAPIServer

	// provided for easier embedding
	Informers internalinformers.SharedInformerFactory
}

// Complete fills in any fields not set that are required to have valid data. It's mutating the receiver.
func (cfg *Config) Complete() CompletedConfig {
	c := completedConfig{
		cfg.GenericConfig.Complete(),
		&cfg.ExtraConfig,
	}

	c.GenericConfig.EnableDiscovery = false
	c.GenericConfig.Version = &version.Info{
		Major: "0",
		Minor: "1",
	}

	return CompletedConfig{&c}
}

// New returns a new instance of CustomResourceDefinitions from the given config.
func (c completedConfig) New(delegationTarget genericapiserver.DelegationTarget) (*CustomResourceDefinitions, error) {
	genericServer, err := c.GenericConfig.New("apiextensions-apiserver", delegationTarget)
	if err != nil {
		return nil, err
	}

	s := &CustomResourceDefinitions{
		GenericAPIServer: genericServer,
	}

	apiResourceConfig := c.GenericConfig.MergedResourceConfig
	apiGroupInfo := genericapiserver.NewDefaultAPIGroupInfo(apiextensions.GroupName, Scheme, metav1.ParameterCodec, Codecs)
	if apiResourceConfig.VersionEnabled(v1beta1.SchemeGroupVersion) {
		storage := map[string]rest.Storage{}
		// customresourcedefinitions
		customResourceDefintionStorage := customresourcedefinition.NewREST(Scheme, c.GenericConfig.RESTOptionsGetter)
		storage["customresourcedefinitions"] = customResourceDefintionStorage
		storage["customresourcedefinitions/status"] = customresourcedefinition.NewStatusREST(Scheme, customResourceDefintionStorage)

		apiGroupInfo.VersionedResourcesStorageMap[v1beta1.SchemeGroupVersion.Version] = storage
	}
	if apiResourceConfig.VersionEnabled(v1.SchemeGroupVersion) {
		storage := map[string]rest.Storage{}
		// customresourcedefinitions
		customResourceDefintionStorage := customresourcedefinition.NewREST(Scheme, c.GenericConfig.RESTOptionsGetter)
		storage["customresourcedefinitions"] = customResourceDefintionStorage
		storage["customresourcedefinitions/status"] = customresourcedefinition.NewStatusREST(Scheme, customResourceDefintionStorage)

		apiGroupInfo.VersionedResourcesStorageMap[v1.SchemeGroupVersion.Version] = storage
	}

	if err := s.GenericAPIServer.InstallAPIGroup(&apiGroupInfo); err != nil {
		return nil, err
	}

	crdClient, err := internalclientset.NewForConfig(s.GenericAPIServer.LoopbackClientConfig)
	if err != nil {
		// it's really bad that this is leaking here, but until we can fix the test (which I'm pretty sure isn't even testing what it wants to test),
		// we need to be able to move forward
		return nil, fmt.Errorf("failed to create clientset: %v", err)
	}
	s.Informers = internalinformers.NewSharedInformerFactory(crdClient, 5*time.Minute)

	delegateHandler := delegationTarget.UnprotectedHandler()
	if delegateHandler == nil {
		delegateHandler = http.NotFoundHandler()
	}

	versionDiscoveryHandler := &versionDiscoveryHandler{
		discovery: map[schema.GroupVersion]*discovery.APIVersionHandler{},
		delegate:  delegateHandler,
	}
	groupDiscoveryHandler := &groupDiscoveryHandler{
		discovery: map[string]*discovery.APIGroupHandler{},
		delegate:  delegateHandler,
	}
	establishingController := establish.NewEstablishingController(s.Informers.Apiextensions().InternalVersion().CustomResourceDefinitions(), crdClient.Apiextensions())
	crdHandler, err := NewCustomResourceDefinitionHandler(
		versionDiscoveryHandler,
		groupDiscoveryHandler,
		s.Informers.Apiextensions().InternalVersion().CustomResourceDefinitions(),
		delegateHandler,
		c.ExtraConfig.CRDRESTOptionsGetter,
		c.GenericConfig.AdmissionControl,
		establishingController,
		c.ExtraConfig.ServiceResolver,
		c.ExtraConfig.AuthResolverWrapper,
		c.ExtraConfig.MasterCount,
		s.GenericAPIServer.Authorizer,
		c.GenericConfig.RequestTimeout,
		time.Duration(c.GenericConfig.MinRequestTimeout)*time.Second,
		apiGroupInfo.StaticOpenAPISpec,
		c.GenericConfig.MaxRequestBodyBytes,
	)
	if err != nil {
		return nil, err
	}
	s.GenericAPIServer.Handler.NonGoRestfulMux.Handle("/apis", crdHandler)
	s.GenericAPIServer.Handler.NonGoRestfulMux.HandlePrefix("/apis/", crdHandler)

	crdController := NewDiscoveryController(s.Informers.Apiextensions().InternalVersion().CustomResourceDefinitions(), versionDiscoveryHandler, groupDiscoveryHandler)
	namingController := status.NewNamingConditionController(s.Informers.Apiextensions().InternalVersion().CustomResourceDefinitions(), crdClient.Apiextensions())
	nonStructuralSchemaController := nonstructuralschema.NewConditionController(s.Informers.Apiextensions().InternalVersion().CustomResourceDefinitions(), crdClient.Apiextensions())
	apiApprovalController := apiapproval.NewKubernetesAPIApprovalPolicyConformantConditionController(s.Informers.Apiextensions().InternalVersion().CustomResourceDefinitions(), crdClient.Apiextensions())
	finalizingController := finalizer.NewCRDFinalizer(
		s.Informers.Apiextensions().InternalVersion().CustomResourceDefinitions(),
		crdClient.Apiextensions(),
		crdHandler,
	)
	var openapiController *openapicontroller.Controller
	if utilfeature.DefaultFeatureGate.Enabled(apiextensionsfeatures.CustomResourcePublishOpenAPI) {
		openapiController = openapicontroller.NewController(s.Informers.Apiextensions().InternalVersion().CustomResourceDefinitions())
	}

	s.GenericAPIServer.AddPostStartHookOrDie("start-apiextensions-informers", func(context genericapiserver.PostStartHookContext) error {
		s.Informers.Start(context.StopCh)
		return nil
	})
	s.GenericAPIServer.AddPostStartHookOrDie("start-apiextensions-controllers", func(context genericapiserver.PostStartHookContext) error {
		// OpenAPIVersionedService and StaticOpenAPISpec are populated in generic apiserver PrepareRun().
		// Together they serve the /openapi/v2 endpoint on a generic apiserver. A generic apiserver may
		// choose to not enable OpenAPI by having null openAPIConfig, and thus OpenAPIVersionedService
		// and StaticOpenAPISpec are both null. In that case we don't run the CRD OpenAPI controller.
		if utilfeature.DefaultFeatureGate.Enabled(apiextensionsfeatures.CustomResourcePublishOpenAPI) && s.GenericAPIServer.OpenAPIVersionedService != nil && s.GenericAPIServer.StaticOpenAPISpec != nil {
			go openapiController.Run(s.GenericAPIServer.StaticOpenAPISpec, s.GenericAPIServer.OpenAPIVersionedService, context.StopCh)
		}

		go crdController.Run(context.StopCh)
		go namingController.Run(context.StopCh)
		go establishingController.Run(context.StopCh)
		go nonStructuralSchemaController.Run(5, context.StopCh)
		go apiApprovalController.Run(5, context.StopCh)
		go finalizingController.Run(5, context.StopCh)
		return nil
	})
	s.GenericAPIServer.AddPostStartHookOrDie("crd-discovery-available", func(context genericapiserver.PostStartHookContext) error {
		return wait.PollImmediateUntil(100*time.Millisecond, func() (bool, error) {
			// only check if we have a valid list for a given resourceversion
			if !s.Informers.Apiextensions().InternalVersion().CustomResourceDefinitions().Informer().HasSynced() {
				return false, nil
			}

			// The returned group and resource lists might be non-nil with partial results even in the
			// case of non-nil error.  If API aggregation fails, we still want our other discovery information because the CRDs
			// may all be present.
			_, serverGroupsAndResources, discoveryErr := crdClient.Discovery().ServerGroupsAndResources()
			if discoveryErr != nil {
				klog.V(2).Info(discoveryErr)
			}

			serverCRDs, err := s.Informers.Apiextensions().InternalVersion().CustomResourceDefinitions().Lister().List(labels.Everything())
			if err != nil {
				return false, err
			}
			crdGroupsAndResources := sets.NewString()
			for _, crd := range serverCRDs {
				// Skip not active CRD
				if !apiextensions.IsCRDConditionTrue(crd, apiextensions.Established) {
					continue
				}
				if skipOpenshiftAPIServerCRDs.Has(fmt.Sprintf("%s.%s", crd.Spec.Names.Plural, crd.Spec.Group)) {
					continue
				}
				for _, version := range crd.Spec.Versions {
					// Skip versions that are not served
					if !version.Served {
						continue
					}
					crdGroupsAndResources.Insert(fmt.Sprintf("%s.%s.%s", crd.Spec.Names.Plural, version.Name, crd.Spec.Group))
				}
			}

			discoveryGroupsAndResources := sets.NewString()
			for _, resourceList := range serverGroupsAndResources {
				for _, apiResource := range resourceList.APIResources {
					group, version := splitGroupVersion(resourceList.GroupVersion)
					discoveryGroupsAndResources.Insert(fmt.Sprintf("%s.%s.%s", apiResource.Name, version, group))
				}
			}
			if !discoveryGroupsAndResources.HasAll(crdGroupsAndResources.List()...) {
				klog.Infof("waiting for CRD resources in discovery: %#v", crdGroupsAndResources.Difference(discoveryGroupsAndResources))
				return false, nil
			}
			return true, nil
		}, context.StopCh)
	})

	// we don't want to report healthy until we can handle all CRDs that have already been registered.  Waiting for the informer
	// to sync makes sure that the lister will be valid before we begin.  There may still be races for CRDs added after startup,
	// but we won't go healthy until we can handle the ones already present.
	s.GenericAPIServer.AddPostStartHookOrDie("crd-informer-synced", func(context genericapiserver.PostStartHookContext) error {
		return wait.PollImmediateUntil(100*time.Millisecond, func() (bool, error) {
			return s.Informers.Apiextensions().InternalVersion().CustomResourceDefinitions().Informer().HasSynced(), nil
		}, context.StopCh)
	})

	return s, nil
}

func splitGroupVersion(gv string) (group string, version string) {
	ss := strings.SplitN(gv, "/", 2)
	if len(ss) == 1 {
		version = ss[0]
	} else {
		group, version = ss[0], ss[1]
	}
	return
}

func DefaultAPIResourceConfigSource() *serverstorage.ResourceConfig {
	ret := serverstorage.NewResourceConfig()
	// NOTE: GroupVersions listed here will be enabled by default. Don't put alpha versions in the list.
	ret.EnableVersions(
		v1beta1.SchemeGroupVersion,
		v1.SchemeGroupVersion,
	)

	return ret
}
<|endoftext|>


<commit_before>package dns

import (
	"regexp"
)

// IPLocalhost is a regex patter for localhost IP address range.
const IPLocalhost = `((127\.([0-9]{1,3}.){2}[0-9]{1,3})|(::1))`

var localhostIPRegexp = regexp.MustCompile(IPLocalhost)

// IsLocalhost returns true if ip matches the localhost IP regular expression.
// Used for determining if nameserver settings are being passed which are
// localhost addresses
func IsLocalhost(ip string) bool {
	return localhostIPRegexp.MatchString(ip)
}
<commit_msg>fix the regexp for matching an ip address<commit_after>package dns

import (
	"regexp"
)

// IPLocalhost is a regex patter for localhost IP address range.
const IPLocalhost = `((127\.([0-9]{1,3}\.){2}[0-9]{1,3})|(::1))`

var localhostIPRegexp = regexp.MustCompile(IPLocalhost)

// IsLocalhost returns true if ip matches the localhost IP regular expression.
// Used for determining if nameserver settings are being passed which are
// localhost addresses
func IsLocalhost(ip string) bool {
	return localhostIPRegexp.MatchString(ip)
}
<|endoftext|>


<commit_before>package connection

import (
	"github.com/davidscholberg/irkbot/lib/configure"
	"github.com/davidscholberg/irkbot/lib/message"
	"github.com/davidscholberg/irkbot/lib/module"
	"github.com/thoj/go-ircevent"
)

func GetIrcConn(cfg *configure.Config) (*irc.Connection, error) {
	conn := irc.IRC(cfg.User.Nick, cfg.User.User)
	conn.UseTLS = cfg.Server.UseTls
	conn.VerboseCallbackHandler = cfg.Connection.VerboseCallbackHandler
	conn.Debug = cfg.Connection.Debug

	conn.AddCallback("001", func(e *irc.Event) {
		if cfg.User.Identify && conn.GetNick() == cfg.User.Nick {
			conn.Privmsgf("NickServ", "identify %s", cfg.User.Password)
		}
		conn.Join(cfg.Channel.ChannelName)
	})

	conn.AddCallback("366", func(e *irc.Event) {
		if len(cfg.Channel.Greeting) != 0 {
			conn.Privmsg(e.Arguments[1], cfg.Channel.Greeting)
		}
	})

	// TODO: start multiple sayLoops, one per conn
	// TODO: pass conn to sayLoop instead of privmsg callbacks?
	outChan := make(chan message.OutboundMsg)
	go message.SayLoop(outChan)

	err := module.RegisterModules(conn, cfg, outChan)
	if err != nil {
		return conn, err
	}

	return conn, nil
}
<commit_msg>added sleep between NickServ identify and channel join<commit_after>package connection

import (
	"github.com/davidscholberg/irkbot/lib/configure"
	"github.com/davidscholberg/irkbot/lib/message"
	"github.com/davidscholberg/irkbot/lib/module"
	"github.com/thoj/go-ircevent"
	"time"
)

func GetIrcConn(cfg *configure.Config) (*irc.Connection, error) {
	conn := irc.IRC(cfg.User.Nick, cfg.User.User)
	conn.UseTLS = cfg.Server.UseTls
	conn.VerboseCallbackHandler = cfg.Connection.VerboseCallbackHandler
	conn.Debug = cfg.Connection.Debug

	conn.AddCallback("001", func(e *irc.Event) {
		if cfg.User.Identify && conn.GetNick() == cfg.User.Nick {
			conn.Privmsgf("NickServ", "identify %s", cfg.User.Password)
			// temporary horrible hack to allow time to be identified
			// before joining a channel
			time.Sleep(time.Second * 10)
		}
		conn.Join(cfg.Channel.ChannelName)
	})

	conn.AddCallback("366", func(e *irc.Event) {
		if len(cfg.Channel.Greeting) != 0 {
			conn.Privmsg(e.Arguments[1], cfg.Channel.Greeting)
		}
	})

	// TODO: start multiple sayLoops, one per conn
	// TODO: pass conn to sayLoop instead of privmsg callbacks?
	outChan := make(chan message.OutboundMsg)
	go message.SayLoop(outChan)

	err := module.RegisterModules(conn, cfg, outChan)
	if err != nil {
		return conn, err
	}

	return conn, nil
}
<|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.

// The following enables go generate to generate the doc.go file.
//go:generate go run $V23_ROOT/release/go/src/v.io/x/lib/cmdline/testdata/gendoc.go . -help

package main

import (
	"encoding/base64"
	"encoding/json"
	"errors"
	"fmt"

	"v.io/v23"
	"v.io/v23/context"
	"v.io/v23/security"
	"v.io/x/lib/cmdline"
	"v.io/x/ref/lib/security/securityflag"
	"v.io/x/ref/lib/signals"
	"v.io/x/ref/lib/v23cmd"
	_ "v.io/x/ref/runtime/factories/generic"
	"v.io/x/ref/services/device/internal/claim"
	"v.io/x/ref/services/identity"
)

var (
	permsDir     string
	blessingRoot string
)

func runServer(ctx *context.T, _ *cmdline.Env, _ []string) error {
	if blessingRoot != "" {
		addRoot(ctx, blessingRoot)
	}

	auth := securityflag.NewAuthorizerOrDie()
	claimable, claimed := claim.NewClaimableDispatcher(ctx, permsDir, "", auth)
	if claimable == nil {
		return errors.New("device is already claimed")
	}

	server, err := v23.NewServer(ctx)
	if err != nil {
		return err
	}
	if _, err := server.Listen(v23.GetListenSpec(ctx)); err != nil {
		return err
	}
	if err := server.ServeDispatcher("", claimable); err != nil {
		return err
	}

	status := server.Status()
	ctx.Infof("Listening on: %v", status.Endpoints)
	if len(status.Endpoints) > 0 {
		fmt.Printf("NAME=%s\n", status.Endpoints[0].Name())
	}
	select {
	case <-claimed:
		return nil
	case s := <-signals.ShutdownOnSignals(ctx):
		return fmt.Errorf("received signal %v", s)
	}
}

func addRoot(ctx *context.T, jRoot string) {
	var bRoot identity.BlessingRootResponse
	if err := json.Unmarshal([]byte(jRoot), &bRoot); err != nil {
		ctx.Fatalf("unable to unmarshal the json blessing root: %v", err)
	}
	decodedKey, err := base64.URLEncoding.DecodeString(bRoot.PublicKey)
	if err != nil {
		ctx.Fatalf("unable to decode public key: %v", err)
	}
	key, err := security.UnmarshalPublicKey(decodedKey)
	if err != nil {
		ctx.Fatalf("unable to unmarshal the public key: %v", err)
	}
	roots := v23.GetPrincipal(ctx).Roots()
	for _, name := range bRoot.Names {
		if err := roots.Add(key, security.BlessingPattern(name)); err != nil {
			ctx.Fatalf("unable to add root: %v", err)
		}
	}
}

func main() {
	rootCmd := &cmdline.Command{
		Name:  "claimable",
		Short: "Run claimable server",
		Long: `
Claimable is a server that implements the Claimable interface from
v.io/v23/services/device. It exits immediately if the device is already
claimed. Otherwise, it keeps running until a successful Claim() request
is received.

It uses -v23.permissions.* to authorize the Claim request.
`,
		Runner: v23cmd.RunnerFunc(runServer),
	}
	rootCmd.Flags.StringVar(&permsDir, "perms-dir", "", "The directory where permissions will be stored.")
	rootCmd.Flags.StringVar(&blessingRoot, "blessing-root", "", "The blessing root to trust, JSON-encoded, e.g. from https://v.io/auth/blessing-root")
	cmdline.Main(rootCmd)
}
<commit_msg>services/device/claimable: use the roaming profile<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.

// The following enables go generate to generate the doc.go file.
//go:generate go run $V23_ROOT/release/go/src/v.io/x/lib/cmdline/testdata/gendoc.go . -help

package main

import (
	"encoding/base64"
	"encoding/json"
	"errors"
	"fmt"

	"v.io/v23"
	"v.io/v23/context"
	"v.io/v23/security"
	"v.io/x/lib/cmdline"
	"v.io/x/ref/lib/security/securityflag"
	"v.io/x/ref/lib/signals"
	"v.io/x/ref/lib/v23cmd"
	_ "v.io/x/ref/runtime/factories/roaming"
	"v.io/x/ref/services/device/internal/claim"
	"v.io/x/ref/services/identity"
)

var (
	permsDir     string
	blessingRoot string
)

func runServer(ctx *context.T, _ *cmdline.Env, _ []string) error {
	if blessingRoot != "" {
		addRoot(ctx, blessingRoot)
	}

	auth := securityflag.NewAuthorizerOrDie()
	claimable, claimed := claim.NewClaimableDispatcher(ctx, permsDir, "", auth)
	if claimable == nil {
		return errors.New("device is already claimed")
	}

	server, err := v23.NewServer(ctx)
	if err != nil {
		return err
	}
	if _, err := server.Listen(v23.GetListenSpec(ctx)); err != nil {
		return err
	}
	if err := server.ServeDispatcher("", claimable); err != nil {
		return err
	}

	status := server.Status()
	ctx.Infof("Listening on: %v", status.Endpoints)
	if len(status.Endpoints) > 0 {
		fmt.Printf("NAME=%s\n", status.Endpoints[0].Name())
	}
	select {
	case <-claimed:
		return nil
	case s := <-signals.ShutdownOnSignals(ctx):
		return fmt.Errorf("received signal %v", s)
	}
}

func addRoot(ctx *context.T, jRoot string) {
	var bRoot identity.BlessingRootResponse
	if err := json.Unmarshal([]byte(jRoot), &bRoot); err != nil {
		ctx.Fatalf("unable to unmarshal the json blessing root: %v", err)
	}
	decodedKey, err := base64.URLEncoding.DecodeString(bRoot.PublicKey)
	if err != nil {
		ctx.Fatalf("unable to decode public key: %v", err)
	}
	key, err := security.UnmarshalPublicKey(decodedKey)
	if err != nil {
		ctx.Fatalf("unable to unmarshal the public key: %v", err)
	}
	roots := v23.GetPrincipal(ctx).Roots()
	for _, name := range bRoot.Names {
		if err := roots.Add(key, security.BlessingPattern(name)); err != nil {
			ctx.Fatalf("unable to add root: %v", err)
		}
	}
}

func main() {
	rootCmd := &cmdline.Command{
		Name:  "claimable",
		Short: "Run claimable server",
		Long: `
Claimable is a server that implements the Claimable interface from
v.io/v23/services/device. It exits immediately if the device is already
claimed. Otherwise, it keeps running until a successful Claim() request
is received.

It uses -v23.permissions.* to authorize the Claim request.
`,
		Runner: v23cmd.RunnerFunc(runServer),
	}
	rootCmd.Flags.StringVar(&permsDir, "perms-dir", "", "The directory where permissions will be stored.")
	rootCmd.Flags.StringVar(&blessingRoot, "blessing-root", "", "The blessing root to trust, JSON-encoded, e.g. from https://v.io/auth/blessing-root")
	cmdline.Main(rootCmd)
}
<|endoftext|>


<commit_before>package captainslog

import (
	"encoding/json"
	"fmt"
	"strings"
)

// JSONKeyTransformer is a Transformer implementation that finds periods in JSON
// keys in CEE syslog messages and replaces them. This can be used in
// conjunction with systems such as Elasticsearch 2.x which do not
// fully support ECMA-404 (for instance, Elasticsearch 2.x does
// not allow periods in key names, which ECMA-404 does)
type JSONKeyTransformer struct {
	old      string
	new      string
	replacer *strings.Replacer
}

// NewJSONKeyTransformer begins construction of a JSONKeyTransformer.
func NewJSONKeyTransformer() *JSONKeyTransformer {
	return &JSONKeyTransformer{}
}

// OldString sets the string that will be replaced in JSON keys
func (t *JSONKeyTransformer) OldString(oldstring string) *JSONKeyTransformer {
	t.old = oldstring
	return t
}

// NewString sets the string that OldString will be converted to
func (t *JSONKeyTransformer) NewString(newstring string) *JSONKeyTransformer {
	t.new = newstring
	return t
}

// Do finishes construction of the JSONKeyTransformer and returns an
// error if any arguments are missing
func (t *JSONKeyTransformer) Do() (*JSONKeyTransformer, error) {
	if t.old == "" || t.new == "" {
		return t, fmt.Errorf("bad arguments")
	}
	t.replacer = strings.NewReplacer(t.old, t.new)
	return t, nil
}

// recurseTransformMap is a helper method to visit multi-level JSON used by Transform
func (t *JSONKeyTransformer) recurseTransformMap(in, out map[string]interface{}) {
	for k, v := range in {
		transformedKey := t.replacer.Replace(k)
		switch cv := v.(type) {
		case map[string]interface{}:
			nv := make(map[string]interface{})
			out[transformedKey] = nv
			t.recurseTransformMap(cv, nv)
		case []interface{}:
			nv := make([]interface{}, len(cv))
			out[transformedKey] = nv
			t.recurseTransformArr(cv, nv)
		default:
			out[transformedKey] = v
		}
	}
}

// recurseTransformArr is a helper method to visit multi-level JSON used by recurseTransformMap
func (t *JSONKeyTransformer) recurseTransformArr(in, out []interface{}) {
	for i, v := range in {
		switch cv := v.(type) {
		case map[string]interface{}:
			nv := make(map[string]interface{})
			out[i] = nv
			t.recurseTransformMap(cv, nv)
		case []interface{}:
			nv := make([]interface{}, len(cv))
			out[i] = nv
			t.recurseTransformArr(cv, nv)
		default:
			out[i] = v
		}
	}
}

// Transform accepts a SyslogMsg, and if it is a CEE syslog message, "fixes"
// the JSON keys to be compatible with Elasticsearch 2.x
func (t *JSONKeyTransformer) Transform(msg SyslogMsg) (SyslogMsg, error) {
	if !msg.IsCee {
		return msg, ErrTransform
	}

	transformedStructured := make(map[string]interface{})
	t.recurseTransformMap(msg.JSONValues, transformedStructured)
	newContent, _ := json.Marshal(transformedStructured)
	msg.Content = string(newContent)
	msg.JSONValues = transformedStructured
	return msg, nil
}
<commit_msg>Problem: used a keyword as a struct member<commit_after>package captainslog

import (
	"encoding/json"
	"fmt"
	"strings"
)

// JSONKeyTransformer is a Transformer implementation that finds periods in JSON
// keys in CEE syslog messages and replaces them. This can be used in
// conjunction with systems such as Elasticsearch 2.x which do not
// fully support ECMA-404 (for instance, Elasticsearch 2.x does
// not allow periods in key names, which ECMA-404 does)
type JSONKeyTransformer struct {
	oldVal   string
	newVal   string
	replacer *strings.Replacer
}

// NewJSONKeyTransformer begins construction of a JSONKeyTransformer.
func NewJSONKeyTransformer() *JSONKeyTransformer {
	return &JSONKeyTransformer{}
}

// OldString sets the string that will be replaced in JSON keys
func (t *JSONKeyTransformer) OldString(oldstring string) *JSONKeyTransformer {
	t.oldVal = oldstring
	return t
}

// NewString sets the string that OldString will be converted to
func (t *JSONKeyTransformer) NewString(newstring string) *JSONKeyTransformer {
	t.newVal = newstring
	return t
}

// Do finishes construction of the JSONKeyTransformer and returns an
// error if any arguments are missing
func (t *JSONKeyTransformer) Do() (*JSONKeyTransformer, error) {
	if t.oldVal == "" || t.newVal == "" {
		return t, fmt.Errorf("bad arguments")
	}
	t.replacer = strings.NewReplacer(t.oldVal, t.newVal)
	return t, nil
}

// recurseTransformMap is a helper method to visit multi-level JSON used by Transform
func (t *JSONKeyTransformer) recurseTransformMap(in, out map[string]interface{}) {
	for k, v := range in {
		transformedKey := t.replacer.Replace(k)
		switch cv := v.(type) {
		case map[string]interface{}:
			nv := make(map[string]interface{})
			out[transformedKey] = nv
			t.recurseTransformMap(cv, nv)
		case []interface{}:
			nv := make([]interface{}, len(cv))
			out[transformedKey] = nv
			t.recurseTransformArr(cv, nv)
		default:
			out[transformedKey] = v
		}
	}
}

// recurseTransformArr is a helper method to visit multi-level JSON used by recurseTransformMap
func (t *JSONKeyTransformer) recurseTransformArr(in, out []interface{}) {
	for i, v := range in {
		switch cv := v.(type) {
		case map[string]interface{}:
			nv := make(map[string]interface{})
			out[i] = nv
			t.recurseTransformMap(cv, nv)
		case []interface{}:
			nv := make([]interface{}, len(cv))
			out[i] = nv
			t.recurseTransformArr(cv, nv)
		default:
			out[i] = v
		}
	}
}

// Transform accepts a SyslogMsg, and if it is a CEE syslog message, "fixes"
// the JSON keys to be compatible with Elasticsearch 2.x
func (t *JSONKeyTransformer) Transform(msg SyslogMsg) (SyslogMsg, error) {
	if !msg.IsCee {
		return msg, ErrTransform
	}

	transformedStructured := make(map[string]interface{})
	t.recurseTransformMap(msg.JSONValues, transformedStructured)
	newContent, _ := json.Marshal(transformedStructured)
	msg.Content = string(newContent)
	msg.JSONValues = transformedStructured
	return msg, nil
}
<|endoftext|>


<commit_before><commit_msg>Remove unused wrappedfs.go<commit_after><|endoftext|>