<commit_before>package coordinator

import (
	"sync/atomic"

	"github.com/privacylab/talek/common"
	//"github.com/privacylab/talek/cuckoo"
	//"golang.org/x/net/trace"
)

// Server is the main logic for the central coordinator
type Server struct {
	/** Private State **/
	// Static
	log  *common.Logger
	name string

	// Thread-safe
	config    atomic.Value  // Config
	commitLog []*CommitArgs // Append and read only

	// Channels
	commitChan chan *CommitArgs
}

// NewServer creates a new Centralized talek server.
func NewServer(name string, config common.Config) *Server {
	s := &Server{}
	s.log = common.NewLogger(name)
	s.name = name
	s.config.Store(config)

	return s
}

/**********************************
 * PUBLIC RPC METHODS (threadsafe)
 **********************************/

// GetInfo returns information about this server
func (s *Server) GetInfo(args *interface{}, reply *GetInfoReply) error {
	reply.Err = ""
	reply.Name = s.name
	return nil
}

// GetCommonConfig returns the common global config
func (s *Server) GetCommonConfig(args *interface{}, reply *common.Config) error {
	config := s.config.Load().(common.Config)
	*reply = config
	return nil
}

// Commit accepts a single Write to commit. The
func (s *Server) Commit(args *CommitArgs, reply *CommitReply) error {
	s.commitChan <- args
	reply.Err = ""
	return nil
}

/**
func (s *Server) GetUpdates(args *common.GetUpdatesArgs, reply *common.GetUpdatesReply) error {
}
**/

/**********************************
 * PUBLIC LOCAL METHODS (threadsafe)
 **********************************/

// Close shuts down the server
func (s *Server) Close() {
}

/**********************************
 * PRIVATE METHODS (single-threaded)
 **********************************/

// processCommits will read from s.commitChan and properly trigger work
func (s *Server) processCommits() {
	var commit *CommitArgs
	//conf := s.config.Load().(common.Config)

	for {
		select {
		case commit = <-s.commitChan:
			s.commitLog = append(s.commitLog, commit)
		}
	}
}
<commit_msg>setup triggers for a layout build<commit_after>package coordinator

import (
	"sync/atomic"
	"time"

	"github.com/privacylab/talek/common"
	//"github.com/privacylab/talek/cuckoo"
	//"golang.org/x/net/trace"
)

// Server is the main logic for the central coordinator
type Server struct {
	/** Private State **/
	// Static
	log            *common.Logger
	name           string
	buildThreshold uint64
	buildInterval  time.Duration

	// Thread-safe
	config        atomic.Value  // Config
	commitLog     []*CommitArgs // Append and read only
	numNewCommits uint64
	timeLastBuild time.Time
	buildCount    uint64

	// Channels
	commitChan chan *CommitArgs
}

// NewServer creates a new Centralized talek server.
func NewServer(name string, config common.Config, buildThreshold uint64, buildInterval time.Duration) *Server {
	s := &Server{}
	s.log = common.NewLogger(name)
	s.name = name
	s.buildThreshold = buildThreshold
	s.buildInterval = buildInterval

	s.config.Store(config)
	s.commitLog = make([]*CommitArgs, 0)
	s.numNewCommits = 0
	s.timeLastBuild = time.Now()
	s.buildCount = 0

	go s.processCommits()
	return s
}

/**********************************
 * PUBLIC RPC METHODS (threadsafe)
 **********************************/

// GetInfo returns information about this server
func (s *Server) GetInfo(args *interface{}, reply *GetInfoReply) error {
	reply.Err = ""
	reply.Name = s.name
	return nil
}

// GetCommonConfig returns the common global config
func (s *Server) GetCommonConfig(args *interface{}, reply *common.Config) error {
	config := s.config.Load().(common.Config)
	*reply = config
	return nil
}

// Commit accepts a single Write to commit. The
func (s *Server) Commit(args *CommitArgs, reply *CommitReply) error {
	s.commitChan <- args
	reply.Err = ""
	return nil
}

/**
func (s *Server) GetUpdates(args *common.GetUpdatesArgs, reply *common.GetUpdatesReply) error {
}
**/

/**********************************
 * PUBLIC LOCAL METHODS (threadsafe)
 **********************************/

// Close shuts down the server
func (s *Server) Close() {
}

/**********************************
 * PRIVATE METHODS (single-threaded)
 **********************************/

// processCommits will read from s.commitChan and properly trigger work
func (s *Server) processCommits() {
	var commit *CommitArgs
	conf := s.config.Load().(common.Config)
	windowSize := conf.WindowSize()
	tick := time.After(s.buildInterval)

	triggerBuild := func() {
		// Garbage collect old items
		idx := 0
		if len(s.commitLog) > windowSize {
			idx = len(s.commitLog) - windowSize
		}
		s.commitLog = s.commitLog[idx:]
		// Spawn build goroutine
		go s.buildLayout(s.buildCount, s.commitLog[:])
		// Reset state
		s.numNewCommits = 0
		s.buildCount += 1
	}

	for {
		select {
		// Handle new commits
		case commit = <-s.commitChan:
			s.commitLog = append(s.commitLog, commit)
			s.numNewCommits += 1
			// Trigger build if over threshold
			if s.numNewCommits > s.buildThreshold {
				triggerBuild()
			}
			continue
		// Periodically trigger a build
		case <-tick:
			triggerBuild()
			tick = time.After(s.buildInterval) // Re-up timer
			continue
		}
	}
}

func (s *Server) buildLayout(buildId uint64, commitLog []*CommitArgs) {
	// Construct cuckoo table layout

	// Construct global interest vector

	// Push layout to replicas

	// Push global interest vector to global frontends

}
<|endoftext|>


<commit_before>package coordinator

import (
	"crypto/rand"
	"encoding/binary"
	"fmt"
	"sync"
	"time"

	"github.com/privacylab/talek/common"
	"github.com/privacylab/talek/cuckoo"
	"golang.org/x/net/trace"
)

// Server is the main logic for the central coordinator
type Server struct {
	/** Private State **/
	// Static
	log               *common.Logger
	name              string
	snapshotThreshold uint64
	snapshotInterval  time.Duration

	// Thread-safe (locked)
	lock          sync.Mutex
	config        common.Config // Config
	servers       []NotifyInterface
	commitLog     []*CommitArgs // Append and read only
	numNewCommits uint64
	snapshotCount uint64
	lastLayout    []uint64
	cuckooData    []byte
	cuckooTable   *cuckoo.Table

	// Channels
	notifyChan chan bool
}

// NewServer creates a new Centralized talek server.
func NewServer(name string, config common.Config, servers []NotifyInterface, snapshotThreshold uint64, snapshotInterval time.Duration) (*Server, error) {
	s := &Server{}
	s.log = common.NewLogger(name)
	s.name = name
	s.snapshotThreshold = snapshotThreshold
	s.snapshotInterval = snapshotInterval

	s.lock = sync.Mutex{}
	s.config = config
	if servers == nil {
		s.servers = make([]NotifyInterface, 0)
	} else {
		s.servers = servers
	}
	s.commitLog = make([]*CommitArgs, 0)
	s.numNewCommits = 0
	s.snapshotCount = 0
	s.lastLayout = nil
	s.cuckooData = make([]byte, config.NumBuckets*config.BucketDepth*uint64(common.IDSize))

	// Choose a random seed for the cuckoo table
	seedBytes := make([]byte, 8)
	_, err := rand.Read(seedBytes)
	if err != nil {
		s.log.Error.Printf("coordinator.NewServer(%v) error: %v", name, err)
		return nil, err
	}
	seed, _ := binary.Varint(seedBytes)
	s.cuckooTable = cuckoo.NewTable(name, config.NumBuckets, config.BucketDepth, config.DataSize, s.cuckooData, seed)
	s.notifyChan = make(chan bool)

	go s.loop()

	s.log.Info.Printf("coordinator.NewServer(%v) success\n", name)
	return s, nil
}

/**********************************
 * PUBLIC RPC METHODS (threadsafe)
 **********************************/

// GetInfo returns information about this server
func (s *Server) GetInfo(args *interface{}, reply *GetInfoReply) error {
	tr := trace.New("Coordinator", "GetInfo")
	defer tr.Finish()
	s.lock.Lock()

	reply.Err = ""
	reply.Name = s.name
	reply.SnapshotID = s.snapshotCount

	s.lock.Unlock()
	return nil
}

// GetCommonConfig returns the common global config
func (s *Server) GetCommonConfig(args *interface{}, reply *common.Config) error {
	tr := trace.New("Coordinator", "GetCommonConfig")
	defer tr.Finish()
	s.lock.Lock()

	*reply = s.config

	s.lock.Unlock()
	return nil
}

// Commit accepts a single Write to commit. The
func (s *Server) Commit(args *CommitArgs, reply *CommitReply) error {
	tr := trace.New("Coordinator", "Commit")
	defer tr.Finish()
	reply.Err = ""

	s.lock.Lock()

	windowSize := s.config.WindowSize()
	// Garbage Collect old elements
	for uint64(len(s.commitLog)) >= windowSize {
		_ = s.cuckooTable.Remove(asCuckooItem(s.commitLog[0]))
		s.commitLog = s.commitLog[1:]
	}

	// Insert new item
	s.numNewCommits++
	s.commitLog = append(s.commitLog, args)
	ok, _ := s.cuckooTable.Insert(asCuckooItem(args))
	if !ok {
		s.log.Error.Fatalf("%v.processCommit failed to insert new element", s.name)
		return fmt.Errorf("Error inserting into cuckoo table")
	}
	//s.log.Info.Printf("%v\n", data)

	s.lock.Unlock()

	// Do notifications in loop()
	s.notifyChan <- false
	return nil
}

/**
func (s *Server) GetUpdates(args *common.GetUpdatesArgs, reply *common.GetUpdatesReply) error {
	tr := trace.New("Coordinator", "GetUpdates")
	defer tr.Finish()
}
**/

/**********************************
 * PUBLIC LOCAL METHODS (threadsafe)
 **********************************/

// Close shuts down the server
func (s *Server) Close() {
	s.log.Info.Printf("%v.Close: success", s.name)
}

func (s *Server) AddServer(server NotifyInterface) {
	s.lock.Lock()
	s.servers = append(s.servers, server)
	s.lock.Unlock()
}

// NotifySnapshot notifies the current cuckoo layout out
// If `force` is false, ignore when under a threshold
// Returns: true if snapshot was built, false if ignored
func (s *Server) NotifySnapshot(force bool) bool {
	s.lock.Lock()

	// Ignore if under threshold and not forcing
	if !force {
		if s.numNewCommits < s.snapshotThreshold {
			s.lock.Unlock()
			return false
		}
	}

	// Reset state
	s.numNewCommits = 0
	s.snapshotCount++

	// Construct global interest vector
	// intVec := buildGlobalInterestVector(commitLog)
	// @todo

	// Copy the layout
	s.lastLayout = make([]uint64, len(s.cuckooData)/8)
	for i := 0; i < len(s.lastLayout); i++ {
		idx := i * 8
		s.lastLayout[i], _ = binary.Uvarint(s.cuckooData[idx:(idx + 8)])
	}

	// Sync with buildGlobalInterestVector goroutine
	// @todo

	// Send notifications
	go sendNotification(s.servers[:], s.snapshotCount)
	s.lock.Unlock()

	return true
}

/**********************************
 * PRIVATE METHODS (single-threaded)
 **********************************/
// Periodically call NotifySnapshot
func (s *Server) loop() {
	tick := time.After(s.snapshotInterval)
	for {
		select {
		// Called after Commit
		case <-s.notifyChan:
			ok := s.NotifySnapshot(false)
			if ok { // Re-up timer if snapshot built
				tick = time.After(s.snapshotInterval)
			}
			continue
		// Periodically trigger a snapshot
		case <-tick:
			s.NotifySnapshot(true)
			tick = time.After(s.snapshotInterval) // Re-up timer
			continue
		}
	}
}

/**********************************
 * HELPER FUNCTIONS
 **********************************/

// Converts a CommitArgs to a cuckoo.Item
func asCuckooItem(args *CommitArgs) *cuckoo.Item {
	itemData := make([]byte, common.IDSize)
	binary.PutUvarint(itemData, args.ID)
	return &cuckoo.Item{
		ID:      args.ID,
		Data:    itemData,
		Bucket1: args.Bucket1,
		Bucket2: args.Bucket2,
	}
}

func buildGlobalInterestVector() {
	// @todo
}

func sendNotification(servers []NotifyInterface, snapshotID uint64) {
	args := &NotifyArgs{
		SnapshotID: snapshotID,
	}
	for _, server := range servers {
		// Ignoring errors and replies
		go server.Notify(args, &NotifyReply{})
	}
}
<commit_msg>log errors from notification<commit_after>package coordinator

import (
	"crypto/rand"
	"encoding/binary"
	"fmt"
	"sync"
	"time"

	"github.com/privacylab/talek/common"
	"github.com/privacylab/talek/cuckoo"
	"golang.org/x/net/trace"
)

// Server is the main logic for the central coordinator
type Server struct {
	/** Private State **/
	// Static
	log               *common.Logger
	name              string
	snapshotThreshold uint64
	snapshotInterval  time.Duration

	// Thread-safe (locked)
	lock          sync.Mutex
	config        common.Config // Config
	servers       []NotifyInterface
	commitLog     []*CommitArgs // Append and read only
	numNewCommits uint64
	snapshotCount uint64
	lastLayout    []uint64
	cuckooData    []byte
	cuckooTable   *cuckoo.Table

	// Channels
	notifyChan chan bool
}

// NewServer creates a new Centralized talek server.
func NewServer(name string, config common.Config, servers []NotifyInterface, snapshotThreshold uint64, snapshotInterval time.Duration) (*Server, error) {
	s := &Server{}
	s.log = common.NewLogger(name)
	s.name = name
	s.snapshotThreshold = snapshotThreshold
	s.snapshotInterval = snapshotInterval

	s.lock = sync.Mutex{}
	s.config = config
	if servers == nil {
		s.servers = make([]NotifyInterface, 0)
	} else {
		s.servers = servers
	}
	s.commitLog = make([]*CommitArgs, 0)
	s.numNewCommits = 0
	s.snapshotCount = 0
	s.lastLayout = nil
	s.cuckooData = make([]byte, config.NumBuckets*config.BucketDepth*uint64(common.IDSize))

	// Choose a random seed for the cuckoo table
	seedBytes := make([]byte, 8)
	_, err := rand.Read(seedBytes)
	if err != nil {
		s.log.Error.Printf("coordinator.NewServer(%v) error: %v", name, err)
		return nil, err
	}
	seed, _ := binary.Varint(seedBytes)
	s.cuckooTable = cuckoo.NewTable(name, config.NumBuckets, config.BucketDepth, config.DataSize, s.cuckooData, seed)
	s.notifyChan = make(chan bool)

	go s.loop()

	s.log.Info.Printf("coordinator.NewServer(%v) success\n", name)
	return s, nil
}

/**********************************
 * PUBLIC RPC METHODS (threadsafe)
 **********************************/

// GetInfo returns information about this server
func (s *Server) GetInfo(args *interface{}, reply *GetInfoReply) error {
	tr := trace.New("Coordinator", "GetInfo")
	defer tr.Finish()
	s.lock.Lock()

	reply.Err = ""
	reply.Name = s.name
	reply.SnapshotID = s.snapshotCount

	s.lock.Unlock()
	return nil
}

// GetCommonConfig returns the common global config
func (s *Server) GetCommonConfig(args *interface{}, reply *common.Config) error {
	tr := trace.New("Coordinator", "GetCommonConfig")
	defer tr.Finish()
	s.lock.Lock()

	*reply = s.config

	s.lock.Unlock()
	return nil
}

// Commit accepts a single Write to commit. The
func (s *Server) Commit(args *CommitArgs, reply *CommitReply) error {
	tr := trace.New("Coordinator", "Commit")
	defer tr.Finish()
	reply.Err = ""

	s.lock.Lock()

	windowSize := s.config.WindowSize()
	// Garbage Collect old elements
	for uint64(len(s.commitLog)) >= windowSize {
		_ = s.cuckooTable.Remove(asCuckooItem(s.commitLog[0]))
		s.commitLog = s.commitLog[1:]
	}

	// Insert new item
	s.numNewCommits++
	s.commitLog = append(s.commitLog, args)
	ok, _ := s.cuckooTable.Insert(asCuckooItem(args))
	if !ok {
		s.log.Error.Fatalf("%v.processCommit failed to insert new element", s.name)
		return fmt.Errorf("Error inserting into cuckoo table")
	}
	//s.log.Info.Printf("%v\n", data)

	s.lock.Unlock()

	// Do notifications in loop()
	s.notifyChan <- false
	return nil
}

/**
func (s *Server) GetUpdates(args *common.GetUpdatesArgs, reply *common.GetUpdatesReply) error {
	tr := trace.New("Coordinator", "GetUpdates")
	defer tr.Finish()
}
**/

/**********************************
 * PUBLIC LOCAL METHODS (threadsafe)
 **********************************/

// Close shuts down the server
func (s *Server) Close() {
	s.log.Info.Printf("%v.Close: success", s.name)
}

func (s *Server) AddServer(server NotifyInterface) {
	s.lock.Lock()
	s.servers = append(s.servers, server)
	s.lock.Unlock()
}

// NotifySnapshot notifies the current cuckoo layout out
// If `force` is false, ignore when under a threshold
// Returns: true if snapshot was built, false if ignored
func (s *Server) NotifySnapshot(force bool) bool {
	s.lock.Lock()

	// Ignore if under threshold and not forcing
	if !force {
		if s.numNewCommits < s.snapshotThreshold {
			s.lock.Unlock()
			return false
		}
	}

	// Reset state
	s.numNewCommits = 0
	s.snapshotCount++

	// Construct global interest vector
	// intVec := buildGlobalInterestVector(commitLog)
	// @todo

	// Copy the layout
	s.lastLayout = make([]uint64, len(s.cuckooData)/8)
	for i := 0; i < len(s.lastLayout); i++ {
		idx := i * 8
		s.lastLayout[i], _ = binary.Uvarint(s.cuckooData[idx:(idx + 8)])
	}

	// Sync with buildGlobalInterestVector goroutine
	// @todo

	// Send notifications
	go sendNotification(s.log, s.servers[:], s.snapshotCount)
	s.lock.Unlock()

	return true
}

/**********************************
 * PRIVATE METHODS (single-threaded)
 **********************************/
// Periodically call NotifySnapshot
func (s *Server) loop() {
	tick := time.After(s.snapshotInterval)
	for {
		select {
		// Called after Commit
		case <-s.notifyChan:
			ok := s.NotifySnapshot(false)
			if ok { // Re-up timer if snapshot built
				tick = time.After(s.snapshotInterval)
			}
			continue
		// Periodically trigger a snapshot
		case <-tick:
			s.NotifySnapshot(true)
			tick = time.After(s.snapshotInterval) // Re-up timer
			continue
		}
	}
}

/**********************************
 * HELPER FUNCTIONS
 **********************************/

// Converts a CommitArgs to a cuckoo.Item
func asCuckooItem(args *CommitArgs) *cuckoo.Item {
	itemData := make([]byte, common.IDSize)
	binary.PutUvarint(itemData, args.ID)
	return &cuckoo.Item{
		ID:      args.ID,
		Data:    itemData,
		Bucket1: args.Bucket1,
		Bucket2: args.Bucket2,
	}
}

func buildGlobalInterestVector() {
	// @todo
}

func sendNotification(log *common.Logger, servers []NotifyInterface, snapshotID uint64) {
	args := &NotifyArgs{
		SnapshotID: snapshotID,
	}
	doNotify := func(s NotifyInterface, args *NotifyArgs) {
		err := s.Notify(args, &NotifyReply{})
		if err != nil {
			log.Error.Printf("sendNotification failed: %v", err)
		}
	}
	for _, s := range servers {
		go doNotify(s, args)
	}
}
<|endoftext|>


<commit_before>// SPDX-License-Identifier: MIT

package web

import (
	"bytes"
	"net/http"
	"net/http/httptest"
	"testing"

	"github.com/issue9/assert"
	"golang.org/x/text/language"
	"golang.org/x/text/message"

	"github.com/issue9/web/result"
)

func TestContext_NewResult(t *testing.T) {
	a := assert.New(t)
	srv := newServer(a)
	srv.AddResultMessage(400, 40000, "lang") // lang 有翻译
	w := httptest.NewRecorder()

	// 能正常翻译错误信息
	r := httptest.NewRequest(http.MethodGet, "/path", nil)
	r.Header.Set("accept-language", language.SimplifiedChinese.String())
	r.Header.Set("accept", "application/json")
	ctx := srv.NewContext(w, r)
	rslt := ctx.NewResult(40000)
	rslt.Render()
	a.Equal(w.Body.String(), `{"message":"hans","code":40000}`)

	// 未指定 accept-language，采用默认的 und
	w = httptest.NewRecorder()
	r = httptest.NewRequest(http.MethodGet, "/path", nil)
	r.Header.Set("accept", "application/json")
	ctx = srv.NewContext(w, r)
	rslt = ctx.NewResult(40000)
	rslt.Render()
	a.Equal(w.Body.String(), `{"message":"und","code":40000}`)

	// 不存在的本地化信息，采用默认的 und
	w = httptest.NewRecorder()
	r = httptest.NewRequest(http.MethodGet, "/path", nil)
	r.Header.Set("accept-language", "en-US")
	r.Header.Set("accept", "application/json")
	ctx = srv.NewContext(w, r)
	rslt = ctx.NewResult(40000)
	rslt.Render()
	a.Equal(w.Body.String(), `{"message":"und","code":40000}`)

	// 不存在
	a.Panic(func() { ctx.NewResult(400) })
	a.Panic(func() { ctx.NewResult(50000) })
}

func TestContext_NewResultWithFields(t *testing.T) {
	a := assert.New(t)

	r := httptest.NewRequest(http.MethodGet, "/path", bytes.NewBufferString("123"))
	r.Header.Set("Accept", "application/json")
	r.Header.Set("Content-Type", "application/json")
	w := httptest.NewRecorder()
	ctx := newServer(a).NewContext(w, r)
	ctx.server.AddResultMessage(http.StatusBadRequest, 40010, "40010")
	ctx.server.AddResultMessage(http.StatusBadRequest, 40011, "40011")

	rslt := ctx.NewResultWithFields(40010, result.Fields{
		"k1": []string{"v1", "v2"},
	})
	a.True(rslt.HasFields())

	rslt.Render()
	a.Equal(w.Body.String(), `{"message":"40010","code":40010,"fields":[{"name":"k1","message":["v1","v2"]}]}`)
}

func TestServer_ResultMessages(t *testing.T) {
	a := assert.New(t)
	srv := newServer(a)
	a.NotNil(srv)

	a.NotPanic(func() {
		srv.AddResultMessage(400, 40010, "lang")
	})

	lmsgs := srv.ResultMessages(message.NewPrinter(language.Und, message.Catalog(srv.catalog)))
	a.Equal(lmsgs[40010], "und")

	lmsgs = srv.ResultMessages(message.NewPrinter(language.SimplifiedChinese, message.Catalog(srv.catalog)))
	a.Equal(lmsgs[40010], "hans")

	lmsgs = srv.ResultMessages(message.NewPrinter(language.TraditionalChinese, message.Catalog(srv.catalog)))
	a.Equal(lmsgs[40010], "hant")

	lmsgs = srv.ResultMessages(message.NewPrinter(language.English, message.Catalog(srv.catalog)))
	a.Equal(lmsgs[40010], "und")
}
<commit_msg>test: 添加针对 SetErrorHandler 和 Result 的相互影响测试<commit_after>// SPDX-License-Identifier: MIT

package web

import (
	"bytes"
	"net/http"
	"net/http/httptest"
	"testing"

	"github.com/issue9/assert"
	"golang.org/x/text/language"
	"golang.org/x/text/message"

	"github.com/issue9/web/result"
)

func TestContext_NewResult(t *testing.T) {
	a := assert.New(t)
	srv := newServer(a)
	srv.SetErrorHandle(func(w http.ResponseWriter, status int) {
		w.WriteHeader(status)
		w.Write([]byte("error-handler"))
	}, 400) // 此处用于检测是否影响 result.Render() 的输出
	srv.AddResultMessage(400, 40000, "lang") // lang 有翻译
	w := httptest.NewRecorder()

	// 能正常翻译错误信息
	r := httptest.NewRequest(http.MethodGet, "/path", nil)
	r.Header.Set("accept-language", language.SimplifiedChinese.String())
	r.Header.Set("accept", "application/json")
	ctx := srv.NewContext(w, r)
	rslt := ctx.NewResult(40000)
	rslt.Render()
	a.Equal(w.Body.String(), `{"message":"hans","code":40000}`)

	// 未指定 accept-language，采用默认的 und
	w = httptest.NewRecorder()
	r = httptest.NewRequest(http.MethodGet, "/path", nil)
	r.Header.Set("accept", "application/json")
	ctx = srv.NewContext(w, r)
	rslt = ctx.NewResult(40000)
	rslt.Render()
	a.Equal(w.Body.String(), `{"message":"und","code":40000}`)

	// 不存在的本地化信息，采用默认的 und
	w = httptest.NewRecorder()
	r = httptest.NewRequest(http.MethodGet, "/path", nil)
	r.Header.Set("accept-language", "en-US")
	r.Header.Set("accept", "application/json")
	ctx = srv.NewContext(w, r)
	rslt = ctx.NewResult(40000)
	rslt.Render()
	a.Equal(w.Body.String(), `{"message":"und","code":40000}`)

	// 不存在
	a.Panic(func() { ctx.NewResult(400) })
	a.Panic(func() { ctx.NewResult(50000) })
}

func TestContext_NewResultWithFields(t *testing.T) {
	a := assert.New(t)

	r := httptest.NewRequest(http.MethodGet, "/path", bytes.NewBufferString("123"))
	r.Header.Set("Accept", "application/json")
	r.Header.Set("Content-Type", "application/json")
	w := httptest.NewRecorder()
	ctx := newServer(a).NewContext(w, r)
	ctx.server.AddResultMessage(http.StatusBadRequest, 40010, "40010")
	ctx.server.AddResultMessage(http.StatusBadRequest, 40011, "40011")

	rslt := ctx.NewResultWithFields(40010, result.Fields{
		"k1": []string{"v1", "v2"},
	})
	a.True(rslt.HasFields())

	rslt.Render()
	a.Equal(w.Body.String(), `{"message":"40010","code":40010,"fields":[{"name":"k1","message":["v1","v2"]}]}`)
}

func TestServer_ResultMessages(t *testing.T) {
	a := assert.New(t)
	srv := newServer(a)
	a.NotNil(srv)

	a.NotPanic(func() {
		srv.AddResultMessage(400, 40010, "lang")
	})

	lmsgs := srv.ResultMessages(message.NewPrinter(language.Und, message.Catalog(srv.catalog)))
	a.Equal(lmsgs[40010], "und")

	lmsgs = srv.ResultMessages(message.NewPrinter(language.SimplifiedChinese, message.Catalog(srv.catalog)))
	a.Equal(lmsgs[40010], "hans")

	lmsgs = srv.ResultMessages(message.NewPrinter(language.TraditionalChinese, message.Catalog(srv.catalog)))
	a.Equal(lmsgs[40010], "hant")

	lmsgs = srv.ResultMessages(message.NewPrinter(language.English, message.Catalog(srv.catalog)))
	a.Equal(lmsgs[40010], "und")
}
<|endoftext|>


<commit_before>package floc

import (
	"fmt"
	"testing"
)

func TestResult_IsNone(t *testing.T) {
	r := None

	if r.IsNone() == false {
		t.Fatalf("%s expects None but has %s", t.Name(), r.String())
	}

	if r.IsValid() == false {
		t.Fatalf("%s expects None to be valid", t.Name())
	}
}

func TestResult_IsCanceled(t *testing.T) {
	r := Canceled

	if r.IsCanceled() == false {
		t.Fatalf("%s expects Canceled but has %s", t.Name(), r.String())
	}

	if r.IsValid() == false {
		t.Fatalf("%s expects Canceled to be valid", t.Name())
	}
}

func TestResult_IsCompleted(t *testing.T) {
	r := Completed

	if r.IsCompleted() == false {
		t.Fatalf("%s expects Completed but has %s", t.Name(), r.String())
	}

	if r.IsValid() == false {
		t.Fatalf("%s expects Completed to be valid", t.Name())
	}
}

func TestResult_IsFailed(t *testing.T) {
	r := Failed

	if r.IsFailed() == false {
		t.Fatalf("%s expects Failed but has %s", t.Name(), r.String())
	}

	if r.IsValid() == false {
		t.Fatalf("%s expects Failed to be valid", t.Name())
	}
}

func TestResult_IsValid(t *testing.T) {
	const n = 1000

	r := Result(n)

	if r.IsValid() == true {
		t.Fatalf("%s expects %s to be invalid", t.Name(), r.String())
	}

	s := fmt.Sprintf("Result(%d)", n)
	if r.String() != s {
		t.Fatalf("%s expects %s but has %s", t.Name(), s, r.String())
	}
}

func TestResult_IsFinished(t *testing.T) {
	if None.IsFinished() == true {
		t.Fatalf("%s expects None to be not finished", t.Name())
	}

	if Completed.IsFinished() == false {
		t.Fatalf("%s expects Completed to be finished", t.Name())
	}

	if Canceled.IsFinished() == false {
		t.Fatalf("%s expects Canceled to be finished", t.Name())
	}

	if Failed.IsFinished() == false {
		t.Fatalf("%s expects Failed to be finished", t.Name())
	}
}

func TestResult_Int32(t *testing.T) {
	var n int32
	for n = 0; n < 1000; n++ {
		r := Result(n)
		if r.i32() != n {
			t.Fatalf("%s expects Result to be %d but has %d", t.Name(), n, r.i32())
		}
	}
}

func TestResult_String(t *testing.T) {
	if None.String() != "None" {
		t.Fatalf("%s expects None bu has %s", t.Name(), None.String())
	}

	if Completed.String() != "Completed" {
		t.Fatalf("%s expects Completed bu has %s", t.Name(), Completed.String())
	}

	if Canceled.String() != "Canceled" {
		t.Fatalf("%s expects Canceled bu has %s", t.Name(), Canceled.String())
	}

	if Failed.String() != "Failed" {
		t.Fatalf("%s expects Failed bu has %s", t.Name(), Failed.String())
	}
}
<commit_msg>Reformat tests<commit_after>package floc

import (
	"fmt"
	"testing"
)

func TestResult_IsNone(t *testing.T) {
	r := None
	if r.IsNone() == false {
		t.Fatalf("%s expects None but has %s", t.Name(), r.String())
	} else if r.IsValid() == false {
		t.Fatalf("%s expects None to be valid", t.Name())
	}
}

func TestResult_IsCanceled(t *testing.T) {
	r := Canceled
	if r.IsCanceled() == false {
		t.Fatalf("%s expects Canceled but has %s", t.Name(), r.String())
	} else if r.IsValid() == false {
		t.Fatalf("%s expects Canceled to be valid", t.Name())
	}
}

func TestResult_IsCompleted(t *testing.T) {
	r := Completed
	if r.IsCompleted() == false {
		t.Fatalf("%s expects Completed but has %s", t.Name(), r.String())
	} else if r.IsValid() == false {
		t.Fatalf("%s expects Completed to be valid", t.Name())
	}
}

func TestResult_IsFailed(t *testing.T) {
	r := Failed
	if r.IsFailed() == false {
		t.Fatalf("%s expects Failed but has %s", t.Name(), r.String())
	} else if r.IsValid() == false {
		t.Fatalf("%s expects Failed to be valid", t.Name())
	}
}

func TestResult_IsValid(t *testing.T) {
	const n = 1000

	r := Result(n)
	if r.IsValid() == true {
		t.Fatalf("%s expects %s to be invalid", t.Name(), r.String())
	}

	s := fmt.Sprintf("Result(%d)", n)
	if r.String() != s {
		t.Fatalf("%s expects %s but has %s", t.Name(), s, r.String())
	}
}

func TestResult_IsFinished(t *testing.T) {
	if None.IsFinished() == true {
		t.Fatalf("%s expects None to be not finished", t.Name())
	}

	if Completed.IsFinished() == false {
		t.Fatalf("%s expects Completed to be finished", t.Name())
	}

	if Canceled.IsFinished() == false {
		t.Fatalf("%s expects Canceled to be finished", t.Name())
	}

	if Failed.IsFinished() == false {
		t.Fatalf("%s expects Failed to be finished", t.Name())
	}
}

func TestResult_Int32(t *testing.T) {
	var n int32
	for n = 0; n < 1000; n++ {
		r := Result(n)
		if r.i32() != n {
			t.Fatalf("%s expects Result to be %d but has %d", t.Name(), n, r.i32())
		}
	}
}

func TestResult_String(t *testing.T) {
	if None.String() != "None" {
		t.Fatalf("%s expects None bu has %s", t.Name(), None.String())
	}

	if Completed.String() != "Completed" {
		t.Fatalf("%s expects Completed bu has %s", t.Name(), Completed.String())
	}

	if Canceled.String() != "Canceled" {
		t.Fatalf("%s expects Canceled bu has %s", t.Name(), Canceled.String())
	}

	if Failed.String() != "Failed" {
		t.Fatalf("%s expects Failed bu has %s", t.Name(), Failed.String())
	}
}
<|endoftext|>


<commit_before>package rhymer_test

import (
    "testing"
    "os"
    "github.com/juanchel/rhymer"
)

var r = rhymer.NewRhymer()

var rhymeTests = []struct {
    a string     // first input
    b string     // second input
    expected int // expected result
} {
    {"cat",    "cat",     1},
    {"do",     "to",      1},
    {"cat",    "bat",     1},
    {"cat",    "acrobat", 1},
    {"over",   "clover",  1},
    {"master", "raster",  1},
    {"masTER", "RaStEr",  1},
    {"aunt",   "rant",    1},
    {"aunt",   "want",    1},
    {"rant",   "want",    0},
    {"do",     "toot",    0},
    {"cat",    "dog",     0},
    {"over",   "ever",    0},
    {"ever",   "clover",  0},
    {"kanye",  "cat",     -1},
    {"kanye",  "yeezy",   -1},
    {"",       "cat",     -1},
    {"",       "",        -1},
    {"^cat",   "&bat",    -1},
}

var phoneticRhymeTests = []struct {
    a string     // first input
    b []string   // second input
    expected int // expected input
} {
    {"cat",   []string{"AE", "T"},            1},
    {"cat",   []string{"S", "AE", "T"},       1},
    {"cat",   []string{"???", "AE", "T"},     1},
    {"hello", []string{"Y", "EH", "L", "OW"}, 1},
    {"cat",   []string{"AE"},                 0},
    {"cat",   []string{"T"},                  0},
    {"cat",   []string{""},                   0},
    {"cat",   []string{"???"},                0},
    {"kanye", []string{"AY"},                 -1},
}

var syllabicReduceTests = []struct {
    a []string        // input
    expected []string // expected
} {
    {[]string{"K", "AE", "T"}, []string{"AE", "T"}},
    {[]string{"K", "AE", "K", "AE", "T"}, []string{"AE", "T"}},
    {[]string{"AE", "T"}, []string{"AE", "T"}},
    {[]string{"T", "T"}, []string{}},
    {[]string{""}, []string{}},
    {[]string{"&&"}, []string{}},
}

func TestMain(m *testing.M) {
    os.Exit(m.Run())
}

func TestRhymes(m *testing.T) {
    for _, i := range rhymeTests {
        actual := r.Rhymes(i.a, i.b)
        if actual != i.expected {
            m.Errorf("Rhymes(%s, %s): expected %d but got %d", i.a, i.b, i.expected, actual)
        }
    }
}

func TestRhymesToPhonetic(m *testing.T) {
    for _, i := range phoneticRhymeTests {
        actual := r.RhymesToPhonetic(i.a, i.b)
        if actual != i.expected {
            m.Errorf("RhymesToPhonetic(%s, %v): expected %d but got %d", i.a, i.b, i.expected, actual)
        }
    }
}

func TestSyllabicReduce(m *testing.T) {
    for _, i := range syllabicReduceTests {
        actual := rhymer.SyllabicReduce(i.a)
        if len(actual) != len(i.expected) {
            m.Errorf("SyllabicReduce(%v) returned the wrong number of phonemes: %v, expected %v", i.a, actual, i.expected)
        } else {
            for n, v := range actual {
                if v != i.expected[n] {
                    m.Errorf("SyllabicReduce(%v) returned the wrong results: %v, expected %v", i.a, actual, i.expected)
                }
            }
        }
    }
}

func TestPronounceSimple(m *testing.T) {
    expected := [3]string{"K", "AE", "T"}
    actual := r.Pronounce("cat")
    if len(actual) != 1 {
        m.Errorf("Pronounce(cat) returned the wrong number of results: %d", len(actual))
    } else if len(actual[0]) != 3 {
        m.Errorf("Pronounce(cat) returned the wrong number of phonemes: %d")
    } else {
        for i, v := range actual[0] {
            if v != expected[i] {
                m.Errorf("Pronounce(cat) returned the wrong phoneme at index %d: got %v, expected %v", i, actual[0], expected)
            }
        }
    }
}

func TestPronounceMultiple(m *testing.T) {
    expectedA := [3]string{"AE", "N", "T"}
    expectedB := [3]string{"AO", "N", "T"}
    actual := r.Pronounce("aunt")
    if len(actual) != 2 {
        m.Errorf("Pronounce(aunt) returned the wrong number of results: %d", len(actual))
    } else if len(actual[0]) != 3 || len(actual[1]) != 3 {
        m.Errorf("Pronounce(aunt) returned the wrong number of phonemes")
    } else {
        for i := range actual[0] {
            if actual[0][i] != expectedA[i] && actual[0][i] != expectedB[i] {
                m.Errorf("Pronounce(aunt) returned the wrong phoneme: got %v, expected [%v %v]", actual, expectedA, expectedB)
            }
            if actual[1][i] != expectedA[i] && actual[1][i] != expectedB[i] {
                m.Errorf("Pronounce(aunt) returned the wrong phoneme: got %v, expected [%v %v]", actual, expectedA, expectedB)
            }
        }
    }
}

func TestPronounceNotFound(m *testing.T) {
    actual := r.Pronounce("naenae")
    if len(actual) != 0 {
        m.Errorf("Pronounce(naenae) should have returned nothing but returned: %v", actual)
    }
    actual = r.Pronounce("!@#$^&")
    if len(actual) != 0 {
        m.Errorf("Pronounce(!@#$^&) should have returned nothing but returned: %v", actual)
    }
}

func TestFindRhymes(m *testing.T) {
    actualWord := r.FindRhymesByWord("crunk")
    actualPhon := r.FindRhymes([]string{"AH", "NG", "K"})
    wordSet := make(map[string]bool)

    for _, v := range actualWord {
        wordSet[v] = true
    }
    for _, v := range actualPhon {
        if !wordSet[v] {
            m.Errorf("Mismatch in FindRhymesByWord(crunk) and FindRhymes([AH NG K])")
        }
    }

    if len(actualWord) != 54 {
        m.Errorf("FindRhymesByWord(crunk) returned %d results, expected 54", len(actualWord))
    }
    if len(actualPhon) != 54 {
        m.Errorf("FindRhymes([AH NG K]) returned %d results, expected 54", len(actualPhon))
    }
}

func TestFindRhymesNotFound(m *testing.T) {
    actualWord := r.FindRhymesByWord("abcd")
    actualPhon := r.FindRhymes([]string{"T", "K", "O"})

    if len(actualWord) != 0 {
        m.Errorf("FindRhymesByWord(abcd) returned %d results, expected 0", len(actualWord))
    }
    if len(actualPhon) != 0 {
        m.Errorf("FindRhymes([AB CD]) returned %d results, expected 0", len(actualPhon))
    }
}<commit_msg>Expand test cases<commit_after>package rhymer_test

import (
    "testing"
    "os"
    "github.com/juanchel/rhymer"
)

var r = rhymer.NewRhymer()

var rhymeTests = []struct {
    a string     // first input
    b string     // second input
    expected int // expected result
} {
    {"cat",    "cat",     1},
    {"do",     "to",      1},
    {"cat",    "bat",     1},
    {"cat",    "acrobat", 1},
    {"over",   "clover",  1},
    {"master", "raster",  1},
    {"masTER", "RaStEr",  1},
    {"aunt",   "rant",    1},
    {"aunt",   "want",    1},
    {"rant",   "want",    0},
    {"do",     "toot",    0},
    {"cat",    "dog",     0},
    {"over",   "ever",    0},
    {"ever",   "clover",  0},
    {"kanye",  "cat",     -1},
    {"kanye",  "yeezy",   -1},
    {"",       "cat",     -1},
    {"",       "",        -1},
    {"^^^",    "&&&ttt",  -1},
    {"好き", "嫌い", -1},
}

var phoneticRhymeTests = []struct {
    a string     // first input
    b []string   // second input
    expected int // expected input
} {
    {"cat",   []string{"AE", "T"},            1},
    {"cat",   []string{"S", "AE", "T"},       1},
    {"cat",   []string{"???", "AE", "T"},     1},
    {"hello", []string{"Y", "EH", "L", "OW"}, 1},
    {"cat",   []string{"AE"},                 0},
    {"cat",   []string{"T"},                  0},
    {"cat",   []string{""},                   0},
    {"cat",   []string{"???"},                0},
    {"",      []string{"???"},                -1},
    {"kanye", []string{"AY"},                 -1},
    {"cat", []string{"再见"}, 0},
    {"kanye", []string{"你好"}, -1},
}

var syllabicReduceTests = []struct {
    a []string        // input
    expected []string // expected
} {
    {[]string{"K", "AE", "T"}, []string{"AE", "T"}},
    {[]string{"K", "AE", "K", "AE", "T"}, []string{"AE", "T"}},
    {[]string{"AE", "T"}, []string{"AE", "T"}},
    {[]string{"T", "T"}, []string{}},
    {[]string{""}, []string{}},
    {[]string{"&&"}, []string{}},
    {[]string{"🔥"}, []string{}},
}

func TestMain(m *testing.M) {
    os.Exit(m.Run())
}

func TestRhymes(m *testing.T) {
    for _, i := range rhymeTests {
        actual := r.Rhymes(i.a, i.b)
        if actual != i.expected {
            m.Errorf("Rhymes(%s, %s): expected %d but got %d", i.a, i.b, i.expected, actual)
        }
    }
}

func TestRhymesToPhonetic(m *testing.T) {
    for _, i := range phoneticRhymeTests {
        actual := r.RhymesToPhonetic(i.a, i.b)
        if actual != i.expected {
            m.Errorf("RhymesToPhonetic(%s, %v): expected %d but got %d", i.a, i.b, i.expected, actual)
        }
    }
}

func TestSyllabicReduce(m *testing.T) {
    for _, i := range syllabicReduceTests {
        actual := rhymer.SyllabicReduce(i.a)
        if len(actual) != len(i.expected) {
            m.Errorf("SyllabicReduce(%v) returned the wrong number of phonemes: %v, expected %v", i.a, actual, i.expected)
        } else {
            for n, v := range actual {
                if v != i.expected[n] {
                    m.Errorf("SyllabicReduce(%v) returned the wrong results: %v, expected %v", i.a, actual, i.expected)
                }
            }
        }
    }
}

func TestPronounceSimple(m *testing.T) {
    expected := [3]string{"K", "AE", "T"}
    actual := r.Pronounce("cat")
    if len(actual) != 1 {
        m.Errorf("Pronounce(cat) returned the wrong number of results: %d", len(actual))
    } else if len(actual[0]) != 3 {
        m.Errorf("Pronounce(cat) returned the wrong number of phonemes: %d")
    } else {
        for i, v := range actual[0] {
            if v != expected[i] {
                m.Errorf("Pronounce(cat) returned the wrong phoneme at index %d: got %v, expected %v", i, actual[0], expected)
            }
        }
    }
}

func TestPronounceMultiple(m *testing.T) {
    expectedA := [3]string{"AE", "N", "T"}
    expectedB := [3]string{"AO", "N", "T"}
    actual := r.Pronounce("aunt")
    if len(actual) != 2 {
        m.Errorf("Pronounce(aunt) returned the wrong number of results: %d", len(actual))
    } else if len(actual[0]) != 3 || len(actual[1]) != 3 {
        m.Errorf("Pronounce(aunt) returned the wrong number of phonemes")
    } else {
        for i := range actual[0] {
            if actual[0][i] != expectedA[i] && actual[0][i] != expectedB[i] {
                m.Errorf("Pronounce(aunt) returned the wrong phoneme: got %v, expected [%v %v]", actual, expectedA, expectedB)
            }
            if actual[1][i] != expectedA[i] && actual[1][i] != expectedB[i] {
                m.Errorf("Pronounce(aunt) returned the wrong phoneme: got %v, expected [%v %v]", actual, expectedA, expectedB)
            }
        }
    }
}

func TestPronounceNotFound(m *testing.T) {
    actual := r.Pronounce("naenae")
    if len(actual) != 0 {
        m.Errorf("Pronounce(naenae) should have returned nothing but returned: %v", actual)
    }
    actual = r.Pronounce("!@#$^&")
    if len(actual) != 0 {
        m.Errorf("Pronounce(!@#$^&) should have returned nothing but returned: %v", actual)
    }
}

func TestFindRhymes(m *testing.T) {
    actualWord := r.FindRhymesByWord("crunk")
    actualPhon := r.FindRhymes([]string{"AH", "NG", "K"})
    wordSet := make(map[string]bool)

    for _, v := range actualWord {
        wordSet[v] = true
    }
    for _, v := range actualPhon {
        if !wordSet[v] {
            m.Errorf("Mismatch in FindRhymesByWord(crunk) and FindRhymes([AH NG K])")
        }
    }

    if len(actualWord) != 54 {
        m.Errorf("FindRhymesByWord(crunk) returned %d results, expected 54", len(actualWord))
    }
    if len(actualPhon) != 54 {
        m.Errorf("FindRhymes([AH NG K]) returned %d results, expected 54", len(actualPhon))
    }
}

func TestFindRhymesNotFound(m *testing.T) {
    actualWord := r.FindRhymesByWord("abcd")
    actualPhon := r.FindRhymes([]string{"T", "K", "O"})

    if len(actualWord) != 0 {
        m.Errorf("FindRhymesByWord(abcd) returned %d results, expected 0", len(actualWord))
    }
    if len(actualPhon) != 0 {
        m.Errorf("FindRhymes([AB CD]) returned %d results, expected 0", len(actualPhon))
    }
}<|endoftext|>


<commit_before>// Copyright 2017 Jeff Foley. All rights reserved.
// Use of this source code is governed by Apache 2 LICENSE that can be found in the LICENSE file.

package sources

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

	"github.com/OWASP/Amass/config"
	eb "github.com/OWASP/Amass/eventbus"
	"github.com/OWASP/Amass/net/http"
	"github.com/OWASP/Amass/requests"
	"github.com/OWASP/Amass/resolvers"
	"github.com/OWASP/Amass/services"
	"github.com/OWASP/Amass/stringset"
)

// Umbrella is the Service that handles access to the Umbrella data source.
type Umbrella struct {
	services.BaseService

	API        *config.APIKey
	SourceType string
	RateLimit  time.Duration
}

// NewUmbrella returns he object initialized, but not yet started.
func NewUmbrella(cfg *config.Config, bus *eb.EventBus, pool *resolvers.ResolverPool) *Umbrella {
	u := &Umbrella{
		SourceType: requests.API,
		RateLimit:  500 * time.Millisecond,
	}

	u.BaseService = *services.NewBaseService(u, "Umbrella", cfg, bus, pool)
	return u
}

// OnStart implements the Service interface
func (u *Umbrella) OnStart() error {
	u.BaseService.OnStart()

	u.API = u.Config().GetAPIKey(u.String())
	if u.API == nil || u.API.Key == "" {
		u.Bus().Publish(requests.LogTopic,
			fmt.Sprintf("%s: API key data was not provided", u.String()),
		)
	}

	go u.processRequests()
	return nil
}

func (u *Umbrella) processRequests() {
	last := time.Now()

	for {
		select {
		case <-u.Quit():
			return
		case req := <-u.DNSRequestChan():
			if u.Config().IsDomainInScope(req.Domain) {
				if time.Now().Sub(last) < u.RateLimit {
					time.Sleep(u.RateLimit)
				}
				last = time.Now()
				u.executeDNSQuery(req.Domain)
				last = time.Now()
			}
		case <-u.AddrRequestChan():
		case <-u.ASNRequestChan():
		case req := <-u.WhoisRequestChan():
			if u.Config().IsDomainInScope(req.Domain) {
				if time.Now().Sub(last) < u.RateLimit {
					time.Sleep(u.RateLimit)
				}
				last = time.Now()
				u.executeWhoisQuery(req.Domain)
				last = time.Now()
			}
		}
	}
}

func (u *Umbrella) executeDNSQuery(domain string) {
	re := u.Config().DomainRegex(domain)
	if re == nil || u.API == nil || u.API.Key == "" {
		return
	}

	u.SetActive()
	headers := u.restHeaders()
	url := u.patternSearchRestURL(domain)
	page, err := http.RequestWebPage(url, nil, headers, "", "")
	if err != nil {
		u.Bus().Publish(requests.LogTopic, fmt.Sprintf("%s: %s: %v", u.String(), url, err))
		return
	}

	for _, name := range re.FindAllString(page, -1) {
		u.Bus().Publish(requests.NewNameTopic, &requests.DNSRequest{
			Name:   cleanName(name),
			Domain: domain,
			Tag:    u.SourceType,
			Source: u.String(),
		})
	}

	url = u.occurrencesRestURL(domain)
	page, err = http.RequestWebPage(url, nil, headers, "", "")
	if err != nil {
		u.Bus().Publish(requests.LogTopic, fmt.Sprintf("%s: %s: %v", u.String(), url, err))
		return
	}

	for _, d := range u.Config().Domains() {
		re := u.Config().DomainRegex(d)
		for _, sd := range re.FindAllString(page, -1) {
			u.Bus().Publish(requests.NewNameTopic, &requests.DNSRequest{
				Name:   cleanName(sd),
				Domain: d,
				Tag:    u.SourceType,
				Source: u.String(),
			})
		}
	}

	u.SetActive()
	url = u.relatedRestURL(domain)
	page, err = http.RequestWebPage(url, nil, headers, "", "")
	if err != nil {
		u.Bus().Publish(requests.LogTopic, fmt.Sprintf("%s: %s: %v", u.String(), url, err))
		return
	}

	for _, d := range u.Config().Domains() {
		re := u.Config().DomainRegex(d)
		for _, sd := range re.FindAllString(page, -1) {
			u.Bus().Publish(requests.NewNameTopic, &requests.DNSRequest{
				Name:   cleanName(sd),
				Domain: d,
				Tag:    u.SourceType,
				Source: u.String(),
			})
		}
	}
}

// Umbrella provides much more than this, but we're only interested in these
// fields
type whoisRecord struct {
	NameServers         []string `json:"nameServers"`
	AdminContactEmail   string   `json:"administrativeContactEmail"`
	BillingContactEmail string   `json:"billingContactEmail"`
	RegistrantEmail     string   `json:"registrantEmail"`
	TechContactEmail    string   `json:"technicalContactEmail"`
	ZoneContactEmail    string   `json:"zoneContactEmail"`
}

// Umbrella provides the same response for email and ns reverse records. Makes
// the json parsing logic simple since we can use the same structs for both
type rWhoisDomain struct {
	Domain  string `json:"domain"`
	Current bool   `json:"current"`
}

type rWhoisResponse struct {
	TotalResults int            `json:"totalResults"`
	MoreData     bool           `json:"moreDataAvailable"`
	Limit        int            `json:"limit"`
	Domains      []rWhoisDomain `json:"domains"`
}

func (u *Umbrella) collateEmails(record *whoisRecord) []string {
	emails := stringset.New()

	if u.validateScope(record.AdminContactEmail) {
		emails.InsertMany(record.AdminContactEmail)
	}
	if u.validateScope(record.BillingContactEmail) {
		emails.InsertMany(record.BillingContactEmail)
	}
	if u.validateScope(record.RegistrantEmail) {
		emails.InsertMany(record.RegistrantEmail)
	}
	if u.validateScope(record.TechContactEmail) {
		emails.InsertMany(record.TechContactEmail)
	}
	if u.validateScope(record.ZoneContactEmail) {
		emails.InsertMany(record.ZoneContactEmail)
	}
	return emails.Slice()
}

func (u *Umbrella) queryWhois(domain string) *whoisRecord {
	var whois whoisRecord
	headers := u.restHeaders()
	whoisURL := u.whoisRecordURL(domain)

	u.SetActive()
	record, err := http.RequestWebPage(whoisURL, nil, headers, "", "")
	if err != nil {
		u.Bus().Publish(requests.LogTopic, fmt.Sprintf("%s: %s: %v", u.String(), whoisURL, err))
		return nil
	}

	err = json.Unmarshal([]byte(record), &whois)
	if err != nil {
		u.Bus().Publish(requests.LogTopic, fmt.Sprintf("%s: %s: %v", u.String(), whoisURL, err))
		return nil
	}

	u.SetActive()
	time.Sleep(u.RateLimit)
	return &whois
}

func (u *Umbrella) queryReverseWhois(apiURL string) []string {
	domains := stringset.New()
	headers := u.restHeaders()
	var whois map[string]rWhoisResponse

	// Umbrella provides data in 500 piece chunks
	for count, more := 0, true; more; count = count + 500 {
		u.SetActive()
		fullAPIURL := fmt.Sprintf("%s&offset=%d", apiURL, count)
		record, err := http.RequestWebPage(fullAPIURL, nil, headers, "", "")
		if err != nil {
			u.Bus().Publish(requests.LogTopic, fmt.Sprintf("%s: %s: %v", u.String(), apiURL, err))
			return domains.Slice()
		}
		err = json.Unmarshal([]byte(record), &whois)

		more = false
		for _, result := range whois {
			if result.TotalResults > 0 {
				for _, domain := range result.Domains {
					if domain.Current {
						domains.Insert(domain.Domain)
					}
				}
			}
			if result.MoreData && more == false {
				more = true
			}
		}

		u.SetActive()
		time.Sleep(u.RateLimit)
	}
	return domains.Slice()
}

func (u *Umbrella) validateScope(input string) bool {
	if input != "" && u.Config().IsDomainInScope(input) {
		return true
	}
	return false
}

func (u *Umbrella) executeWhoisQuery(domain string) {
	if u.API == nil || u.API.Key == "" {
		return
	}

	whoisRecord := u.queryWhois(domain)
	if whoisRecord == nil {
		return
	}

	domains := stringset.New()
	emails := u.collateEmails(whoisRecord)
	if len(emails) > 0 {
		emailURL := u.reverseWhoisByEmailURL(emails...)
		for _, d := range u.queryReverseWhois(emailURL) {
			if !u.Config().IsDomainInScope(d) {
				domains.Insert(d)
			}
		}
	}

	var nameservers []string
	for _, ns := range whoisRecord.NameServers {
		if u.validateScope(ns) {
			nameservers = append(nameservers, ns)
		}
	}
	if len(nameservers) > 0 {
		nsURL := u.reverseWhoisByNSURL(nameservers...)
		for _, d := range u.queryReverseWhois(nsURL) {
			if !u.Config().IsDomainInScope(d) {
				domains.Insert(d)
			}
		}
	}

	if len(domains) > 0 {
		u.Bus().Publish(requests.NewWhoisTopic, &requests.WhoisRequest{
			Domain:     domain,
			NewDomains: domains.Slice(),
			Tag:        u.SourceType,
			Source:     u.String(),
		})
	}
}

func (u *Umbrella) restHeaders() map[string]string {
	headers := map[string]string{"Content-Type": "application/json"}

	if u.API != nil && u.API.Key != "" {
		headers["Authorization"] = "Bearer " + u.API.Key
	}
	return headers

}

func (u *Umbrella) whoisBaseURL() string {
	return `https://investigate.api.umbrella.com/whois/`
}

func (u *Umbrella) whoisRecordURL(domain string) string {
	return u.whoisBaseURL() + domain
}

func (u *Umbrella) reverseWhoisByNSURL(ns ...string) string {
	nameservers := strings.Join(ns, ",")

	return u.whoisBaseURL() + `nameservers?nameServerList=` + nameservers
}

func (u *Umbrella) reverseWhoisByEmailURL(emails ...string) string {
	emailQuery := strings.Join(emails, ",")

	return u.whoisBaseURL() + `emails?emailList=` + emailQuery
}

func (u *Umbrella) patternSearchRestURL(domain string) string {
	return `https://investigate.api.umbrella.com/search/.*[.]` + domain + "?start=-30days&limit=1000"
}

func (u *Umbrella) occurrencesRestURL(domain string) string {
	return "https://investigate.api.umbrella.com/recommendations/name/" + domain + ".json"
}

func (u *Umbrella) relatedRestURL(domain string) string {
	return "https://investigate.api.umbrella.com/links/name/" + domain + ".json"
}
<commit_msg>added support to perform AS and IP addresses information queries<commit_after>// Copyright 2017 Jeff Foley. All rights reserved.
// Use of this source code is governed by Apache 2 LICENSE that can be found in the LICENSE file.

package sources

import (
	"encoding/json"
	"fmt"
	"net"
	"strings"
	"time"

	"github.com/OWASP/Amass/config"
	eb "github.com/OWASP/Amass/eventbus"
	"github.com/OWASP/Amass/net/http"
	"github.com/OWASP/Amass/requests"
	"github.com/OWASP/Amass/resolvers"
	"github.com/OWASP/Amass/services"
	sf "github.com/OWASP/Amass/stringfilter"
	"github.com/OWASP/Amass/stringset"
)

// Umbrella is the Service that handles access to the Umbrella data source.
type Umbrella struct {
	services.BaseService

	API        *config.APIKey
	SourceType string
	RateLimit  time.Duration
	filter     *sf.StringFilter
}

// NewUmbrella returns he object initialized, but not yet started.
func NewUmbrella(cfg *config.Config, bus *eb.EventBus, pool *resolvers.ResolverPool) *Umbrella {
	u := &Umbrella{
		SourceType: requests.API,
		RateLimit:  500 * time.Millisecond,
		filter:     sf.NewStringFilter(),
	}

	u.BaseService = *services.NewBaseService(u, "Umbrella", cfg, bus, pool)
	return u
}

// OnStart implements the Service interface
func (u *Umbrella) OnStart() error {
	u.BaseService.OnStart()

	u.API = u.Config().GetAPIKey(u.String())
	if u.API == nil || u.API.Key == "" {
		u.Bus().Publish(requests.LogTopic,
			fmt.Sprintf("%s: API key data was not provided", u.String()),
		)
	}

	u.Bus().Subscribe(requests.NewAddrTopic, u.SendAddrRequest)
	u.Bus().Subscribe(requests.IPToASNTopic, u.SendASNRequest)
	go u.processRequests()
	return nil
}

func (u *Umbrella) processRequests() {
	last := time.Now()

loop:
	for {
		select {
		case <-u.Quit():
			return
		case req := <-u.DNSRequestChan():
			if u.Config().IsDomainInScope(req.Domain) {
				if time.Now().Sub(last) < u.RateLimit {
					time.Sleep(u.RateLimit)
				}
				last = time.Now()
				u.executeDNSQuery(req.Domain)
				last = time.Now()
			}
		case req := <-u.AddrRequestChan():
			if time.Now().Sub(last) < u.RateLimit {
				time.Sleep(u.RateLimit)
			}
			last = time.Now()
			u.executeAddrQuery(req.Address)
			last = time.Now()
		case req := <-u.ASNRequestChan():
			if req.Address == "" && req.ASN == 0 {
				continue loop
			}
			if time.Now().Sub(last) < u.RateLimit {
				time.Sleep(u.RateLimit)
			}
			last = time.Now()
			if req.Address != "" {
				u.executeASNAddrQuery(req)
			} else {
				u.executeASNQuery(req)
			}
			last = time.Now()
		case req := <-u.WhoisRequestChan():
			if u.Config().IsDomainInScope(req.Domain) {
				if time.Now().Sub(last) < u.RateLimit {
					time.Sleep(u.RateLimit)
				}
				last = time.Now()
				u.executeWhoisQuery(req.Domain)
				last = time.Now()
			}
		}
	}
}

func (u *Umbrella) executeDNSQuery(domain string) {
	if u.API == nil || u.API.Key == "" {
		return
	}

	u.SetActive()
	headers := u.restHeaders()
	url := u.restDNSURL(domain)
	page, err := http.RequestWebPage(url, nil, headers, "", "")
	if err != nil {
		u.Bus().Publish(requests.LogTopic, fmt.Sprintf("%s: %s: %v", u.String(), url, err))
		return
	}
	// Extract the subdomain names from the REST API results
	var subs struct {
		Matches []struct {
			Name string `json:"name"`
		} `json:"matches"`
	}
	if err := json.Unmarshal([]byte(page), &subs); err != nil {
		return
	}

	for _, m := range subs.Matches {
		if d := u.Config().WhichDomain(m.Name); d != "" {
			u.Bus().Publish(requests.NewNameTopic, &requests.DNSRequest{
				Name:   m.Name,
				Domain: d,
				Tag:    u.SourceType,
				Source: u.String(),
			})
		}
	}
}

func (u *Umbrella) executeAddrQuery(addr string) {
	if u.API == nil || u.API.Key == "" {
		return
	}
	if addr == "" || u.filter.Duplicate(addr) {
		return
	}

	u.SetActive()
	headers := u.restHeaders()
	url := u.restAddrURL(addr)
	page, err := http.RequestWebPage(url, nil, headers, "", "")
	if err != nil {
		u.Bus().Publish(requests.LogTopic, fmt.Sprintf("%s: %s: %v", u.String(), url, err))
		return
	}
	// Extract the subdomain names from the REST API results
	var ip struct {
		Records []struct {
			Data string `json:"rr"`
		} `json:"records"`
	}
	if err := json.Unmarshal([]byte(page), &ip); err != nil {
		return
	}

	for _, record := range ip.Records {
		if name := resolvers.RemoveLastDot(record.Data); name != "" {
			if domain := u.Config().WhichDomain(name); domain != "" {
				u.Bus().Publish(requests.NewNameTopic, &requests.DNSRequest{
					Name:   name,
					Domain: domain,
					Tag:    u.SourceType,
					Source: u.String(),
				})
			}
		}
	}
}

func (u *Umbrella) executeASNAddrQuery(req *requests.ASNRequest) {
	if u.API == nil || u.API.Key == "" {
		return
	}

	u.SetActive()
	headers := u.restHeaders()
	url := u.restAddrToASNURL(req.Address)
	page, err := http.RequestWebPage(url, nil, headers, "", "")
	if err != nil {
		u.Bus().Publish(requests.LogTopic, fmt.Sprintf("%s: %s: %v", u.String(), url, err))
		return
	}
	// Extract the AS information from the REST API results
	var as []struct {
		Date        string `json:"creation_date"`
		Registry    int    `json:"ir"`
		Description string `json:"description"`
		ASN         int    `json:"asn"`
		CIDR        string `json:"cidr"`
	}
	if err := json.Unmarshal([]byte(page), &as); err != nil || len(as) == 0 {
		return
	}

	created, err := time.Parse("2006-01-02", as[0].Date)
	if err != nil {
		return
	}

	var registry string
	switch as[0].Registry {
	case 1:
		registry = "AfriNIC"
	case 2:
		registry = "APNIC"
	case 3:
		registry = "ARIN"
	case 4:
		registry = "LACNIC"
	case 5:
		registry = "RIPE NCC"
	default:
		registry = "N/A"
	}

	req.ASN = as[0].ASN
	req.Prefix = as[0].CIDR
	req.Registry = registry
	req.AllocationDate = created
	req.Description = as[0].Description
	req.Tag = u.SourceType
	req.Source = u.String()
	if req.Netblocks == nil {
		req.Netblocks = stringset.New()
		req.Netblocks.Insert(strings.TrimSpace(req.Prefix))
		time.Sleep(u.RateLimit)
		u.executeASNQuery(req)
	}
	u.Bus().Publish(requests.NewASNTopic, req)
}

func (u *Umbrella) executeASNQuery(req *requests.ASNRequest) {
	if u.API == nil || u.API.Key == "" {
		return
	}

	u.SetActive()
	headers := u.restHeaders()
	url := u.restASNToCIDRsURL(req.ASN)
	page, err := http.RequestWebPage(url, nil, headers, "", "")
	if err != nil {
		u.Bus().Publish(requests.LogTopic, fmt.Sprintf("%s: %s: %v", u.String(), url, err))
		return
	}
	// Extract the netblock information from the REST API results
	var netblock []struct {
		CIDR string `json:"cidr"`
		Geo  struct {
			CountryName string `json:"country_name"`
			CountryCode string `json:"country_code"`
		} `json:"geo"`
	}
	if err := json.Unmarshal([]byte(page), &netblock); err != nil || len(netblock) == 0 {
		return
	}

	if req.Netblocks == nil {
		req.Netblocks = stringset.New()
	}

	for _, nb := range netblock {
		req.Netblocks.Insert(strings.TrimSpace(nb.CIDR))
		if nb.CIDR == req.Prefix {
			req.CC = nb.Geo.CountryCode
		}
	}
	// If no basic AS info exists, then obtain an IP and query
	if req.Prefix == "" {
		addr, _, err := net.ParseCIDR(netblock[0].CIDR)

		if err == nil {
			req.Address = addr.String()
			req.CC = netblock[0].Geo.CountryCode
			time.Sleep(u.RateLimit)
			u.executeASNAddrQuery(req)
			return
		}
	}
	// Finish populating the AS info in the request
	for _, nb := range netblock {
		if nb.CIDR == req.Prefix {
			req.CC = nb.Geo.CountryCode
			break
		}
	}
}

// Umbrella provides much more than this, but we're only interested in these
// fields
type whoisRecord struct {
	NameServers         []string `json:"nameServers"`
	AdminContactEmail   string   `json:"administrativeContactEmail"`
	BillingContactEmail string   `json:"billingContactEmail"`
	RegistrantEmail     string   `json:"registrantEmail"`
	TechContactEmail    string   `json:"technicalContactEmail"`
	ZoneContactEmail    string   `json:"zoneContactEmail"`
}

// Umbrella provides the same response for email and ns reverse records. Makes
// the json parsing logic simple since we can use the same structs for both
type rWhoisDomain struct {
	Domain  string `json:"domain"`
	Current bool   `json:"current"`
}

type rWhoisResponse struct {
	TotalResults int            `json:"totalResults"`
	MoreData     bool           `json:"moreDataAvailable"`
	Limit        int            `json:"limit"`
	Domains      []rWhoisDomain `json:"domains"`
}

func (u *Umbrella) collateEmails(record *whoisRecord) []string {
	emails := stringset.New()

	if u.validateScope(record.AdminContactEmail) {
		emails.InsertMany(record.AdminContactEmail)
	}
	if u.validateScope(record.BillingContactEmail) {
		emails.InsertMany(record.BillingContactEmail)
	}
	if u.validateScope(record.RegistrantEmail) {
		emails.InsertMany(record.RegistrantEmail)
	}
	if u.validateScope(record.TechContactEmail) {
		emails.InsertMany(record.TechContactEmail)
	}
	if u.validateScope(record.ZoneContactEmail) {
		emails.InsertMany(record.ZoneContactEmail)
	}
	return emails.Slice()
}

func (u *Umbrella) queryWhois(domain string) *whoisRecord {
	var whois whoisRecord
	headers := u.restHeaders()
	whoisURL := u.whoisRecordURL(domain)

	u.SetActive()
	record, err := http.RequestWebPage(whoisURL, nil, headers, "", "")
	if err != nil {
		u.Bus().Publish(requests.LogTopic, fmt.Sprintf("%s: %s: %v", u.String(), whoisURL, err))
		return nil
	}

	err = json.Unmarshal([]byte(record), &whois)
	if err != nil {
		u.Bus().Publish(requests.LogTopic, fmt.Sprintf("%s: %s: %v", u.String(), whoisURL, err))
		return nil
	}

	u.SetActive()
	time.Sleep(u.RateLimit)
	return &whois
}

func (u *Umbrella) queryReverseWhois(apiURL string) []string {
	domains := stringset.New()
	headers := u.restHeaders()
	var whois map[string]rWhoisResponse

	// Umbrella provides data in 500 piece chunks
	for count, more := 0, true; more; count = count + 500 {
		u.SetActive()
		fullAPIURL := fmt.Sprintf("%s&offset=%d", apiURL, count)
		record, err := http.RequestWebPage(fullAPIURL, nil, headers, "", "")
		if err != nil {
			u.Bus().Publish(requests.LogTopic, fmt.Sprintf("%s: %s: %v", u.String(), apiURL, err))
			return domains.Slice()
		}
		err = json.Unmarshal([]byte(record), &whois)

		more = false
		for _, result := range whois {
			if result.TotalResults > 0 {
				for _, domain := range result.Domains {
					if domain.Current {
						domains.Insert(domain.Domain)
					}
				}
			}
			if result.MoreData && more == false {
				more = true
			}
		}

		u.SetActive()
		time.Sleep(u.RateLimit)
	}
	return domains.Slice()
}

func (u *Umbrella) validateScope(input string) bool {
	if input != "" && u.Config().IsDomainInScope(input) {
		return true
	}
	return false
}

func (u *Umbrella) executeWhoisQuery(domain string) {
	if u.API == nil || u.API.Key == "" {
		return
	}

	whoisRecord := u.queryWhois(domain)
	if whoisRecord == nil {
		return
	}

	domains := stringset.New()
	emails := u.collateEmails(whoisRecord)
	if len(emails) > 0 {
		emailURL := u.reverseWhoisByEmailURL(emails...)
		for _, d := range u.queryReverseWhois(emailURL) {
			if !u.Config().IsDomainInScope(d) {
				domains.Insert(d)
			}
		}
	}

	var nameservers []string
	for _, ns := range whoisRecord.NameServers {
		if u.validateScope(ns) {
			nameservers = append(nameservers, ns)
		}
	}
	if len(nameservers) > 0 {
		nsURL := u.reverseWhoisByNSURL(nameservers...)
		for _, d := range u.queryReverseWhois(nsURL) {
			if !u.Config().IsDomainInScope(d) {
				domains.Insert(d)
			}
		}
	}

	if len(domains) > 0 {
		u.Bus().Publish(requests.NewWhoisTopic, &requests.WhoisRequest{
			Domain:     domain,
			NewDomains: domains.Slice(),
			Tag:        u.SourceType,
			Source:     u.String(),
		})
	}
}

func (u *Umbrella) restHeaders() map[string]string {
	headers := map[string]string{"Content-Type": "application/json"}

	if u.API != nil && u.API.Key != "" {
		headers["Authorization"] = "Bearer " + u.API.Key
	}
	return headers

}

func (u *Umbrella) whoisBaseURL() string {
	return `https://investigate.api.umbrella.com/whois/`
}

func (u *Umbrella) whoisRecordURL(domain string) string {
	return u.whoisBaseURL() + domain
}

func (u *Umbrella) reverseWhoisByNSURL(ns ...string) string {
	nameservers := strings.Join(ns, ",")

	return u.whoisBaseURL() + `nameservers?nameServerList=` + nameservers
}

func (u *Umbrella) reverseWhoisByEmailURL(emails ...string) string {
	emailQuery := strings.Join(emails, ",")

	return u.whoisBaseURL() + `emails?emailList=` + emailQuery
}

func (u *Umbrella) restDNSURL(domain string) string {
	return `https://investigate.api.umbrella.com/search/.*[.]` + domain + "?start=-30days&limit=1000"
}

func (u *Umbrella) restAddrURL(addr string) string {
	return "https://investigate.api.umbrella.com/pdns/ip/" + addr + "?recordType=A,AAAA"
}

func (u *Umbrella) restAddrToASNURL(addr string) string {
	return fmt.Sprintf("https://investigate.api.umbrella.com/bgp_routes/ip/%s/as_for_ip.json", addr)
}

func (u *Umbrella) restASNToCIDRsURL(asn int) string {
	return fmt.Sprintf("https://investigate.api.umbrella.com/bgp_routes/asn/%d/prefixes_for_asn.json", asn)
}
<|endoftext|>


<commit_before>// package merkledag implements the ipfs Merkle DAG datastructures.
package merkledag

import (
	"fmt"

	blocks "github.com/ipfs/go-ipfs/blocks"
	key "github.com/ipfs/go-ipfs/blocks/key"
	bserv "github.com/ipfs/go-ipfs/blockservice"
	"gx/ipfs/QmZy2y8t9zQH2a1b8q2ZSLKp17ATuJoCNxxyMFG5qFExpt/go-net/context"
	logging "gx/ipfs/Qmazh5oNUVsDZTs2g59rq8aYQqwpss8tcUWQzor5sCCEuH/go-log"
)

var log = logging.Logger("merkledag")
var ErrNotFound = fmt.Errorf("merkledag: not found")

// DAGService is an IPFS Merkle DAG service.
type DAGService interface {
	Add(*Node) (key.Key, error)
	AddRecursive(*Node) error
	Get(context.Context, key.Key) (*Node, error)
	Remove(*Node) error
	RemoveRecursive(*Node) error

	// GetDAG returns, in order, all the single leve child
	// nodes of the passed in node.
	GetMany(context.Context, []key.Key) (<-chan *Node, <-chan error)

	Batch() *Batch
}

func NewDAGService(bs *bserv.BlockService) DAGService {
	return &dagService{bs}
}

// dagService is an IPFS Merkle DAG service.
// - the root is virtual (like a forest)
// - stores nodes' data in a BlockService
// TODO: should cache Nodes that are in memory, and be
//       able to free some of them when vm pressure is high
type dagService struct {
	Blocks *bserv.BlockService
}

// Add adds a node to the dagService, storing the block in the BlockService
func (n *dagService) Add(nd *Node) (key.Key, error) {
	if n == nil { // FIXME remove this assertion. protect with constructor invariant
		return "", fmt.Errorf("dagService is nil")
	}

	d, err := nd.Encoded(false)
	if err != nil {
		return "", err
	}

	b := new(blocks.Block)
	b.Data = d
	b.Multihash, err = nd.Multihash()
	if err != nil {
		return "", err
	}

	return n.Blocks.AddBlock(b)
}

func (n *dagService) Batch() *Batch {
	return &Batch{ds: n, MaxSize: 8 * 1024 * 1024}
}

// AddRecursive adds the given node and all child nodes to the BlockService
func (n *dagService) AddRecursive(nd *Node) error {
	_, err := n.Add(nd)
	if err != nil {
		log.Info("AddRecursive Error: %s\n", err)
		return err
	}

	for _, link := range nd.Links {
		if link.Node != nil {
			err := n.AddRecursive(link.Node)
			if err != nil {
				return err
			}
		}
	}

	return nil
}

// Get retrieves a node from the dagService, fetching the block in the BlockService
func (n *dagService) Get(ctx context.Context, k key.Key) (*Node, error) {
	if n == nil {
		return nil, fmt.Errorf("dagService is nil")
	}
	ctx, cancel := context.WithCancel(ctx)
	defer cancel()

	b, err := n.Blocks.GetBlock(ctx, k)
	if err != nil {
		if err == bserv.ErrNotFound {
			return nil, ErrNotFound
		}
		return nil, err
	}

	return Decoded(b.Data)
}

// Remove deletes the given node and all of its children from the BlockService
func (n *dagService) RemoveRecursive(nd *Node) error {
	for _, l := range nd.Links {
		if l.Node != nil {
			n.RemoveRecursive(l.Node)
		}
	}
	k, err := nd.Key()
	if err != nil {
		return err
	}
	return n.Blocks.DeleteBlock(k)
}

func (n *dagService) Remove(nd *Node) error {
	k, err := nd.Key()
	if err != nil {
		return err
	}
	return n.Blocks.DeleteBlock(k)
}

// FetchGraph fetches all nodes that are children of the given node
func FetchGraph(ctx context.Context, root *Node, serv DAGService) error {
	return EnumerateChildrenAsync(ctx, serv, root, key.NewKeySet())
}

// FindLinks searches this nodes links for the given key,
// returns the indexes of any links pointing to it
func FindLinks(links []key.Key, k key.Key, start int) []int {
	var out []int
	for i, lnk_k := range links[start:] {
		if k == lnk_k {
			out = append(out, i+start)
		}
	}
	return out
}

func (ds *dagService) GetMany(ctx context.Context, keys []key.Key) (<-chan *Node, <-chan error) {
	out := make(chan *Node, len(keys))
	errs := make(chan error, 1)
	blocks := ds.Blocks.GetBlocks(ctx, keys)
	var count int

	go func() {
		defer close(out)
		for {
			select {
			case b, ok := <-blocks:
				if !ok {
					if count != len(keys) {
						errs <- fmt.Errorf("failed to fetch all nodes")
					}
					return
				}
				nd, err := Decoded(b.Data)
				if err != nil {
					errs <- err
					return
				}

				// buffered, no need to select
				out <- nd
				count++

			case <-ctx.Done():
				errs <- ctx.Err()
				return
			}
		}
	}()
	return out, errs
}

// GetDAG will fill out all of the links of the given Node.
// It returns a channel of nodes, which the caller can receive
// all the child nodes of 'root' on, in proper order.
func GetDAG(ctx context.Context, ds DAGService, root *Node) []NodeGetter {
	var keys []key.Key
	for _, lnk := range root.Links {
		keys = append(keys, key.Key(lnk.Hash))
	}

	return GetNodes(ctx, ds, keys)
}

// GetNodes returns an array of 'NodeGetter' promises, with each corresponding
// to the key with the same index as the passed in keys
func GetNodes(ctx context.Context, ds DAGService, keys []key.Key) []NodeGetter {

	// Early out if no work to do
	if len(keys) == 0 {
		return nil
	}

	promises := make([]NodeGetter, len(keys))
	sendChans := make([]chan<- *Node, len(keys))
	for i := range keys {
		promises[i], sendChans[i] = newNodePromise(ctx)
	}

	dedupedKeys := dedupeKeys(keys)
	go func() {
		ctx, cancel := context.WithCancel(ctx)
		defer cancel()

		nodechan, errchan := ds.GetMany(ctx, dedupedKeys)

		for count := 0; count < len(keys); {
			select {
			case nd, ok := <-nodechan:
				if !ok {
					return
				}

				k, err := nd.Key()
				if err != nil {
					log.Error("Failed to get node key: ", err)
					continue
				}

				is := FindLinks(keys, k, 0)
				for _, i := range is {
					count++
					sendChans[i] <- nd
				}
			case err := <-errchan:
				log.Error("error fetching: ", err)
				return
			case <-ctx.Done():
				return
			}
		}
	}()
	return promises
}

// Remove duplicates from a list of keys
func dedupeKeys(ks []key.Key) []key.Key {
	kmap := make(map[key.Key]struct{})
	var out []key.Key
	for _, k := range ks {
		if _, ok := kmap[k]; !ok {
			kmap[k] = struct{}{}
			out = append(out, k)
		}
	}
	return out
}

func newNodePromise(ctx context.Context) (NodeGetter, chan<- *Node) {
	ch := make(chan *Node, 1)
	return &nodePromise{
		recv: ch,
		ctx:  ctx,
	}, ch
}

type nodePromise struct {
	cache *Node
	recv  <-chan *Node
	ctx   context.Context
}

// NodeGetter provides a promise like interface for a dag Node
// the first call to Get will block until the Node is received
// from its internal channels, subsequent calls will return the
// cached node.
type NodeGetter interface {
	Get(context.Context) (*Node, error)
}

func (np *nodePromise) Get(ctx context.Context) (*Node, error) {
	if np.cache != nil {
		return np.cache, nil
	}

	select {
	case blk := <-np.recv:
		np.cache = blk
	case <-np.ctx.Done():
		return nil, np.ctx.Err()
	case <-ctx.Done():
		return nil, ctx.Err()
	}
	return np.cache, nil
}

type Batch struct {
	ds *dagService

	blocks  []*blocks.Block
	size    int
	MaxSize int
}

func (t *Batch) Add(nd *Node) (key.Key, error) {
	d, err := nd.Encoded(false)
	if err != nil {
		return "", err
	}

	b := new(blocks.Block)
	b.Data = d
	b.Multihash, err = nd.Multihash()
	if err != nil {
		return "", err
	}

	k := key.Key(b.Multihash)

	t.blocks = append(t.blocks, b)
	t.size += len(b.Data)
	if t.size > t.MaxSize {
		return k, t.Commit()
	}
	return k, nil
}

func (t *Batch) Commit() error {
	_, err := t.ds.Blocks.AddBlocks(t.blocks)
	t.blocks = nil
	t.size = 0
	return err
}

// EnumerateChildren will walk the dag below the given root node and add all
// unseen children to the passed in set.
// TODO: parallelize to avoid disk latency perf hits?
func EnumerateChildren(ctx context.Context, ds DAGService, root *Node, set key.KeySet) error {
	for _, lnk := range root.Links {
		k := key.Key(lnk.Hash)
		if !set.Has(k) {
			set.Add(k)
			child, err := ds.Get(ctx, k)
			if err != nil {
				return err
			}
			err = EnumerateChildren(ctx, ds, child, set)
			if err != nil {
				return err
			}
		}
	}
	return nil
}

func EnumerateChildrenAsync(ctx context.Context, ds DAGService, root *Node, set key.KeySet) error {
	toprocess := make(chan []key.Key, 8)
	nodes := make(chan *Node, 8)
	errs := make(chan error, 1)

	ctx, cancel := context.WithCancel(ctx)
	defer cancel()
	defer close(toprocess)

	go fetchNodes(ctx, ds, toprocess, nodes, errs)

	nodes <- root
	live := 1

	for {
		select {
		case nd, ok := <-nodes:
			if !ok {
				return nil
			}
			// a node has been fetched
			live--

			var keys []key.Key
			for _, lnk := range nd.Links {
				k := key.Key(lnk.Hash)
				if !set.Has(k) {
					set.Add(k)
					live++
					keys = append(keys, k)
				}
			}

			if live == 0 {
				return nil
			}

			if len(keys) > 0 {
				select {
				case toprocess <- keys:
				case <-ctx.Done():
					return ctx.Err()
				}
			}
		case err := <-errs:
			return err
		case <-ctx.Done():
			return ctx.Err()
		}
	}
}

func fetchNodes(ctx context.Context, ds DAGService, in <-chan []key.Key, out chan<- *Node, errs chan<- error) {
	defer close(out)

	get := func(ks []key.Key) {
		nodes, errch := ds.GetMany(ctx, ks)
		for {
			select {
			case nd, ok := <-nodes:
				if !ok {
					return
				}
				select {
				case out <- nd:
				case <-ctx.Done():
					return
				}
			case err := <-errch:
				errs <- err
				return
			}
		}
	}

	for ks := range in {
		go get(ks)
	}
}
<commit_msg>use an option type to simplify concurrency<commit_after>// package merkledag implements the ipfs Merkle DAG datastructures.
package merkledag

import (
	"fmt"
	"sync"

	blocks "github.com/ipfs/go-ipfs/blocks"
	key "github.com/ipfs/go-ipfs/blocks/key"
	bserv "github.com/ipfs/go-ipfs/blockservice"
	"gx/ipfs/QmZy2y8t9zQH2a1b8q2ZSLKp17ATuJoCNxxyMFG5qFExpt/go-net/context"
	logging "gx/ipfs/Qmazh5oNUVsDZTs2g59rq8aYQqwpss8tcUWQzor5sCCEuH/go-log"
)

var log = logging.Logger("merkledag")
var ErrNotFound = fmt.Errorf("merkledag: not found")

// DAGService is an IPFS Merkle DAG service.
type DAGService interface {
	Add(*Node) (key.Key, error)
	AddRecursive(*Node) error
	Get(context.Context, key.Key) (*Node, error)
	Remove(*Node) error
	RemoveRecursive(*Node) error

	// GetDAG returns, in order, all the single leve child
	// nodes of the passed in node.
	GetMany(context.Context, []key.Key) <-chan *NodeOption

	Batch() *Batch
}

func NewDAGService(bs *bserv.BlockService) DAGService {
	return &dagService{bs}
}

// dagService is an IPFS Merkle DAG service.
// - the root is virtual (like a forest)
// - stores nodes' data in a BlockService
// TODO: should cache Nodes that are in memory, and be
//       able to free some of them when vm pressure is high
type dagService struct {
	Blocks *bserv.BlockService
}

// Add adds a node to the dagService, storing the block in the BlockService
func (n *dagService) Add(nd *Node) (key.Key, error) {
	if n == nil { // FIXME remove this assertion. protect with constructor invariant
		return "", fmt.Errorf("dagService is nil")
	}

	d, err := nd.Encoded(false)
	if err != nil {
		return "", err
	}

	b := new(blocks.Block)
	b.Data = d
	b.Multihash, err = nd.Multihash()
	if err != nil {
		return "", err
	}

	return n.Blocks.AddBlock(b)
}

func (n *dagService) Batch() *Batch {
	return &Batch{ds: n, MaxSize: 8 * 1024 * 1024}
}

// AddRecursive adds the given node and all child nodes to the BlockService
func (n *dagService) AddRecursive(nd *Node) error {
	_, err := n.Add(nd)
	if err != nil {
		log.Info("AddRecursive Error: %s\n", err)
		return err
	}

	for _, link := range nd.Links {
		if link.Node != nil {
			err := n.AddRecursive(link.Node)
			if err != nil {
				return err
			}
		}
	}

	return nil
}

// Get retrieves a node from the dagService, fetching the block in the BlockService
func (n *dagService) Get(ctx context.Context, k key.Key) (*Node, error) {
	if n == nil {
		return nil, fmt.Errorf("dagService is nil")
	}
	ctx, cancel := context.WithCancel(ctx)
	defer cancel()

	b, err := n.Blocks.GetBlock(ctx, k)
	if err != nil {
		if err == bserv.ErrNotFound {
			return nil, ErrNotFound
		}
		return nil, err
	}

	return Decoded(b.Data)
}

// Remove deletes the given node and all of its children from the BlockService
func (n *dagService) RemoveRecursive(nd *Node) error {
	for _, l := range nd.Links {
		if l.Node != nil {
			n.RemoveRecursive(l.Node)
		}
	}
	k, err := nd.Key()
	if err != nil {
		return err
	}
	return n.Blocks.DeleteBlock(k)
}

func (n *dagService) Remove(nd *Node) error {
	k, err := nd.Key()
	if err != nil {
		return err
	}
	return n.Blocks.DeleteBlock(k)
}

// FetchGraph fetches all nodes that are children of the given node
func FetchGraph(ctx context.Context, root *Node, serv DAGService) error {
	return EnumerateChildrenAsync(ctx, serv, root, key.NewKeySet())
}

// FindLinks searches this nodes links for the given key,
// returns the indexes of any links pointing to it
func FindLinks(links []key.Key, k key.Key, start int) []int {
	var out []int
	for i, lnk_k := range links[start:] {
		if k == lnk_k {
			out = append(out, i+start)
		}
	}
	return out
}

type NodeOption struct {
	Node *Node
	Err  error
}

func (ds *dagService) GetMany(ctx context.Context, keys []key.Key) <-chan *NodeOption {
	out := make(chan *NodeOption, len(keys))
	blocks := ds.Blocks.GetBlocks(ctx, keys)
	var count int

	go func() {
		defer close(out)
		for {
			select {
			case b, ok := <-blocks:
				if !ok {
					if count != len(keys) {
						out <- &NodeOption{Err: fmt.Errorf("failed to fetch all nodes")}
					}
					return
				}
				nd, err := Decoded(b.Data)
				if err != nil {
					out <- &NodeOption{Err: err}
					return
				}

				// buffered, no need to select
				out <- &NodeOption{Node: nd}
				count++

			case <-ctx.Done():
				out <- &NodeOption{Err: ctx.Err()}
				return
			}
		}
	}()
	return out
}

// GetDAG will fill out all of the links of the given Node.
// It returns a channel of nodes, which the caller can receive
// all the child nodes of 'root' on, in proper order.
func GetDAG(ctx context.Context, ds DAGService, root *Node) []NodeGetter {
	var keys []key.Key
	for _, lnk := range root.Links {
		keys = append(keys, key.Key(lnk.Hash))
	}

	return GetNodes(ctx, ds, keys)
}

// GetNodes returns an array of 'NodeGetter' promises, with each corresponding
// to the key with the same index as the passed in keys
func GetNodes(ctx context.Context, ds DAGService, keys []key.Key) []NodeGetter {

	// Early out if no work to do
	if len(keys) == 0 {
		return nil
	}

	promises := make([]NodeGetter, len(keys))
	sendChans := make([]chan<- *Node, len(keys))
	for i := range keys {
		promises[i], sendChans[i] = newNodePromise(ctx)
	}

	dedupedKeys := dedupeKeys(keys)
	go func() {
		ctx, cancel := context.WithCancel(ctx)
		defer cancel()

		nodechan := ds.GetMany(ctx, dedupedKeys)

		for count := 0; count < len(keys); {
			select {
			case opt, ok := <-nodechan:
				if !ok {
					return
				}

				if opt.Err != nil {
					log.Error("error fetching: ", opt.Err)
					return
				}

				nd := opt.Node

				k, err := nd.Key()
				if err != nil {
					log.Error("Failed to get node key: ", err)
					continue
				}

				is := FindLinks(keys, k, 0)
				for _, i := range is {
					count++
					sendChans[i] <- nd
				}
			case <-ctx.Done():
				return
			}
		}
	}()
	return promises
}

// Remove duplicates from a list of keys
func dedupeKeys(ks []key.Key) []key.Key {
	kmap := make(map[key.Key]struct{})
	var out []key.Key
	for _, k := range ks {
		if _, ok := kmap[k]; !ok {
			kmap[k] = struct{}{}
			out = append(out, k)
		}
	}
	return out
}

func newNodePromise(ctx context.Context) (NodeGetter, chan<- *Node) {
	ch := make(chan *Node, 1)
	return &nodePromise{
		recv: ch,
		ctx:  ctx,
	}, ch
}

type nodePromise struct {
	cache *Node
	recv  <-chan *Node
	ctx   context.Context
}

// NodeGetter provides a promise like interface for a dag Node
// the first call to Get will block until the Node is received
// from its internal channels, subsequent calls will return the
// cached node.
type NodeGetter interface {
	Get(context.Context) (*Node, error)
}

func (np *nodePromise) Get(ctx context.Context) (*Node, error) {
	if np.cache != nil {
		return np.cache, nil
	}

	select {
	case blk := <-np.recv:
		np.cache = blk
	case <-np.ctx.Done():
		return nil, np.ctx.Err()
	case <-ctx.Done():
		return nil, ctx.Err()
	}
	return np.cache, nil
}

type Batch struct {
	ds *dagService

	blocks  []*blocks.Block
	size    int
	MaxSize int
}

func (t *Batch) Add(nd *Node) (key.Key, error) {
	d, err := nd.Encoded(false)
	if err != nil {
		return "", err
	}

	b := new(blocks.Block)
	b.Data = d
	b.Multihash, err = nd.Multihash()
	if err != nil {
		return "", err
	}

	k := key.Key(b.Multihash)

	t.blocks = append(t.blocks, b)
	t.size += len(b.Data)
	if t.size > t.MaxSize {
		return k, t.Commit()
	}
	return k, nil
}

func (t *Batch) Commit() error {
	_, err := t.ds.Blocks.AddBlocks(t.blocks)
	t.blocks = nil
	t.size = 0
	return err
}

// EnumerateChildren will walk the dag below the given root node and add all
// unseen children to the passed in set.
// TODO: parallelize to avoid disk latency perf hits?
func EnumerateChildren(ctx context.Context, ds DAGService, root *Node, set key.KeySet) error {
	for _, lnk := range root.Links {
		k := key.Key(lnk.Hash)
		if !set.Has(k) {
			set.Add(k)
			child, err := ds.Get(ctx, k)
			if err != nil {
				return err
			}
			err = EnumerateChildren(ctx, ds, child, set)
			if err != nil {
				return err
			}
		}
	}
	return nil
}

func EnumerateChildrenAsync(ctx context.Context, ds DAGService, root *Node, set key.KeySet) error {
	toprocess := make(chan []key.Key, 8)
	nodes := make(chan *NodeOption, 8)

	ctx, cancel := context.WithCancel(ctx)
	defer cancel()
	defer close(toprocess)

	go fetchNodes(ctx, ds, toprocess, nodes)

	nodes <- &NodeOption{Node: root}
	live := 1

	for {
		select {
		case opt, ok := <-nodes:
			if !ok {
				return nil
			}

			if opt.Err != nil {
				return opt.Err
			}

			nd := opt.Node

			// a node has been fetched
			live--

			var keys []key.Key
			for _, lnk := range nd.Links {
				k := key.Key(lnk.Hash)
				if !set.Has(k) {
					set.Add(k)
					live++
					keys = append(keys, k)
				}
			}

			if live == 0 {
				return nil
			}

			if len(keys) > 0 {
				select {
				case toprocess <- keys:
				case <-ctx.Done():
					return ctx.Err()
				}
			}
		case <-ctx.Done():
			return ctx.Err()
		}
	}
}

func fetchNodes(ctx context.Context, ds DAGService, in <-chan []key.Key, out chan<- *NodeOption) {
	var wg sync.WaitGroup
	defer func() {
		// wait for all 'get' calls to complete so we don't accidentally send
		// on a closed channel
		wg.Wait()
		close(out)
	}()

	get := func(ks []key.Key) {
		defer wg.Done()
		nodes := ds.GetMany(ctx, ks)
		for opt := range nodes {
			select {
			case out <- opt:
			case <-ctx.Done():
				return
			}
		}
	}

	for ks := range in {
		wg.Add(1)
		go get(ks)
	}
}
<|endoftext|>


<commit_before>package message

import (
	"encoding/json"
	"net/http"
	"net/http/httptest"
	"strings"
	"testing"
	"time"

	"github.com/gorilla/websocket"
)

func makeWsProto(s string) string {
	return "ws" + strings.TrimPrefix(s, "http")
}

func TestIDCounter(t *testing.T) {
	id := IDCounter{}
	if id.Generate() <= 0 {
		t.Errorf("invalid id")
	}
	id.value = 4294967295
	if id.Generate() == 0 {
		t.Errorf("id zero should not be generated")
	}
}

func TestSend(t *testing.T) {
	testSend(t, 1)
}

func BenchmarkSend(b *testing.B) {
	testSend(b, b.N)
}

func testSend(tb testing.TB, n int) {
	done := make(chan struct{})
	var upgrader = websocket.Upgrader{}
	s := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		ws, err := upgrader.Upgrade(w, r, nil)
		if err != nil {
			ws.Close()
			tb.Error(err)
		}
		go readLoop(ws)
		sender := NewSender(ws)

		for i := 0; i < n; i++ {
			sender.Send(OutgoingMessage{
				Ping: &Empty{ID: uint32(i)},
			})
		}
		select {
		case <-done:
		case <-time.After(100 * time.Second):
			tb.Errorf("test timeout")
		}
		sender.Stop()
		select {
		case <-sender.stop:
		case <-time.After(10 * time.Second):
			tb.Errorf("sender was not stopped withing timeout")
		}
	}))
	defer s.Close()

	conn, _, err := websocket.DefaultDialer.Dial(makeWsProto(s.URL), nil)
	if err != nil {
		tb.Error(err)
	}
	defer conn.Close()
	for {
		messageType, r, err := conn.NextReader()
		if err != nil {
			tb.Error(err)
		}
		if messageType != websocket.TextMessage {
			tb.Errorf("incoming message is not a text message")
		}
		var msgs []OutgoingMessage
		e := json.NewDecoder(r)
		if err := e.Decode(&msgs); err != nil {
			tb.Errorf("error unmarshaling message: %v", err)
		}
		if msgs[0].Ping == nil {
			tb.Errorf("ping message not received")
		}
		if msgs[0].Ping.ID >= uint32(n-1) {
			close(done)
			return
		}
	}
}

func readLoop(c *websocket.Conn) {
	for {
		if _, _, err := c.NextReader(); err != nil {
			return
		}
	}
}
<commit_msg>Update Sender test<commit_after>package message

import (
	"net/http"
	"net/http/httptest"
	"strings"
	"testing"
	"time"

	"github.com/gorilla/websocket"
)

func makeWsProto(s string) string {
	return "ws" + strings.TrimPrefix(s, "http")
}

func TestIDCounter(t *testing.T) {
	id := IDCounter{}
	if id.Generate() <= 0 {
		t.Errorf("invalid id")
	}
	id.value = 4294967295
	if id.Generate() == 0 {
		t.Errorf("id zero should not be generated")
	}
}

func TestSend(t *testing.T) {
	testSend(t, 1)
}

func testSend(tb testing.TB, n int) {
	start := make(chan struct{})
	done := make(chan struct{})
	var upgrader = websocket.Upgrader{}
	s := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		ws, err := upgrader.Upgrade(w, r, nil)
		if err != nil {
			ws.Close()
			tb.Error(err)
		}
		go readLoop(ws)
		sender := NewSender(ws)

		<-start
		for i := 0; i < n; i++ {
			sender.Send(OutgoingMessage{
				Ping: &Empty{ID: uint32(i)},
			})
		}
		select {
		case <-done:
		case <-time.After(10 * time.Second):
			tb.Errorf("test timeout")
		}
		sender.Stop()
		select {
		case <-sender.stop:
		case <-time.After(10 * time.Second):
			tb.Errorf("sender was not stopped withing timeout")
		}
	}))
	defer s.Close()

	conn, _, err := websocket.DefaultDialer.Dial(makeWsProto(s.URL), nil)
	if err != nil {
		tb.Error(err)
	}
	defer conn.Close()
	close(start)
	for {
		var msgs []OutgoingMessage
		err := conn.ReadJSON(&msgs)
		if err != nil {
			tb.Error(err)
			continue
		}
		if len(msgs) < 1 {
			tb.Errorf("empty message list received")

		}
		if msgs[0].Ping == nil {
			tb.Errorf("ping message not received")
		}
		if msgs[0].Ping.ID >= uint32(n-1) {
			close(done)
			return
		}
	}
}

func readLoop(c *websocket.Conn) {
	for {
		if _, _, err := c.NextReader(); err != nil {
			return
		}
	}
}
<|endoftext|>


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

// mount_gcsfuse is a small helper for using gcsfuse with mount(8).
//
// mount_gcsfuse can be invoked using a command-line of the form expected for
// mount helpers. It calls the gcsfuse binary, which must be in $PATH, and
// waits for it to complete. The device is passed as --bucket, and other known
// options are converted to appropriate flags.
//
// mount_gcsfuse does not daemonize, and therefore must be used with a wrapper
// that performs daemonization if it is to be used directly with mount(8).
package main

import (
	"errors"
	"log"
	"os"
)

// A 'name=value' mount option. If '=value' is not present, only the name will
// be filled in.
type Option struct {
	Name  string
	Value string
}

// Attempt to parse the terrible undocumented format that mount(8) gives us.
// Return the 'device' (aka 'special' on OS X), the mount point, and a list of
// mount options encountered.
func parseArgs() (device string, mountPoint string, opts []Option, err error) {
	// Example invocation on OS X:
	//
	//     mount -t porp -o key_file=/some\ file.json bucket ~/tmp/mp
	//
	// becomes the following arguments:
	//
	//     Arg 0: "/path/to/mount_gcsfuse"
	//     Arg 1: "-o"
	//     Arg 2: "key_file=/some file.json"
	//     Arg 3: "bucket"
	//     Arg 4: "/Users/jacobsa/tmp/mp"
	//
	// On Linux, the fstab entry
	//
	//     bucket /path/to/mp porp user,key_file=/some\040file.json
	//
	// becomes
	//
	//     Arg 0: "/path/to/mount_gcsfuse"
	//     Arg 1: "bucket"
	//     Arg 2: "/path/to/mp"
	//     Arg 3: "-o"
	//     Arg 4: "rw,noexec,nosuid,nodev,user,key_file=/some file.json"
	//

	err = errors.New("TODO: parseArgs")
	return
}

func main() {
	// Print out each argument.
	//
	// TODO(jacobsa): Get rid of some or all of the debug logging.
	for i, arg := range os.Args {
		log.Printf("Arg %d: %q", i, arg)
	}

	// Attempt to parse arguments.
	device, mountPoint, opts, err := parseArgs()
	if err != nil {
		log.Fatalf("parseArgs: %v", err)
		return
	}

	// Print what we gleaned.
	log.Printf("Device: %q", device)
	log.Printf("Mount point: %q", mountPoint)
	for _, opt := range opts {
		log.Printf("Option %q: %q", opt.Name, opt.Value)
	}

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

// mount_gcsfuse is a small helper for using gcsfuse with mount(8).
//
// mount_gcsfuse can be invoked using a command-line of the form expected for
// mount helpers. It calls the gcsfuse binary, which must be in $PATH, and
// waits for it to complete. The device is passed as --bucket, and other known
// options are converted to appropriate flags.
//
// mount_gcsfuse does not daemonize, and therefore must be used with a wrapper
// that performs daemonization if it is to be used directly with mount(8).
package main

import (
	"errors"
	"fmt"
	"log"
	"os"
)

// A 'name=value' mount option. If '=value' is not present, only the name will
// be filled in.
type Option struct {
	Name  string
	Value string
}

// Parse a single comma-separated list of mount options.
func parseOpts(s string) (opts []Option, err error) {
	err = errors.New("TODO: parseOpts")
	return
}

// Attempt to parse the terrible undocumented format that mount(8) gives us.
// Return the 'device' (aka 'special' on OS X), the mount point, and a list of
// mount options encountered.
func parseArgs() (device string, mountPoint string, opts []Option, err error) {
	// Example invocation on OS X:
	//
	//     mount -t porp -o key_file=/some\ file.json -o ro,blah bucket ~/tmp/mp
	//
	// becomes the following arguments:
	//
	//     Arg 0: "/path/to/mount_gcsfuse"
	//     Arg 1: "-o"
	//     Arg 2: "key_file=/some file.json"
	//     Arg 3: "-o"
	//     Arg 4: "ro"
	//     Arg 5: "-o"
	//     Arg 6: "blah"
	//     Arg 7: "bucket"
	//     Arg 8: "/path/to/mp"
	//
	// On Linux, the fstab entry
	//
	//     bucket /path/to/mp porp user,key_file=/some\040file.json
	//
	// becomes
	//
	//     Arg 0: "/path/to/mount_gcsfuse"
	//     Arg 1: "bucket"
	//     Arg 2: "/path/to/mp"
	//     Arg 3: "-o"
	//     Arg 4: "rw,noexec,nosuid,nodev,user,key_file=/some file.json"
	//

	// Linux and OS X differ on the position of the options. So scan all
	// arguments (aside from the name of the binary), and:
	//
	//  *  Treat the first argument not following "-o" as the device name.
	//  *  Treat the second argument not following "-o" as the mount point.
	//  *  Treat the third argument not following "-o" as an error.
	//  *  Treat all arguments following "-o" as comma-separated options lists.
	//
	rawArgs := 0
	for i, arg := range os.Args[1:] {
		// Skip "-o"; we will look back on the next iteration.
		if arg == "-o" {
			continue
		}

		// If the previous argument was "-o", this is a list of options.
		if os.Args[i-1] == "-o" {
			var tmp []Option
			tmp, err = parseOpts(arg)
			if err != nil {
				err = fmt.Errorf("parseOpts(%q): %v", arg, err)
				return
			}

			opts = append(opts, tmp...)
			continue
		}

		// Otherwise, have we found too many arguments?
		if rawArgs > 2 {
			break
		}

		switch rawArgs {
		case 0:
			device = arg

		case 1:
			mountPoint = arg

		default:
			break
		}

		rawArgs++
	}

	// Did we see all of the raw arguments we expected?
	if rawArgs != 2 {
		err = fmt.Errorf("Expected 2 non-option arguments; got %d", rawArgs)
	}

	err = errors.New("TODO: parseArgs")
	return
}

func main() {
	// Print out each argument.
	//
	// TODO(jacobsa): Get rid of some or all of the debug logging.
	for i, arg := range os.Args {
		log.Printf("Arg %d: %q", i, arg)
	}

	// Attempt to parse arguments.
	device, mountPoint, opts, err := parseArgs()
	if err != nil {
		log.Fatalf("parseArgs: %v", err)
		return
	}

	// Print what we gleaned.
	log.Printf("Device: %q", device)
	log.Printf("Mount point: %q", mountPoint)
	for _, opt := range opts {
		log.Printf("Option %q: %q", opt.Name, opt.Value)
	}

	os.Exit(1)
}
<|endoftext|>