<commit_before>package brokers_test

import (
	"errors"
	"github.com/RichardKnop/machinery/v1"
	"github.com/RichardKnop/machinery/v1/brokers"
	"github.com/RichardKnop/machinery/v1/config"
	"github.com/RichardKnop/machinery/v1/retry"
	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/service/sqs"
	"github.com/stretchr/testify/assert"
	"testing"
	"time"
)

var (
	testBroker           brokers.Interface
	testAWSSQSBroker     *brokers.AWSSQSBroker
	errAWSSQSBroker      *brokers.AWSSQSBroker
	cnf                  *config.Config
	receiveMessageOutput *sqs.ReceiveMessageOutput
)

func init() {
	testAWSSQSBroker = brokers.TestAWSSQSBroker
	errAWSSQSBroker = brokers.ErrAWSSQSBroker
	cnf = brokers.TestConf
	receiveMessageOutput = brokers.ReceiveMessageOutput
	testBroker = brokers.NewAWSSQSBroker(cnf)
}

func TestNewAWSSQSBroker(t *testing.T) {
	assert.IsType(t, testAWSSQSBroker, testBroker)
}

func TestPrivateFunc_continueReceivingMessages(t *testing.T) {
	qURL := testAWSSQSBroker.DefaultQueueURLForTest()
	deliveries := make(chan *sqs.ReceiveMessageOutput)
	firstStep := make(chan int)
	nextStep := make(chan int)
	go func() {
		stopReceivingChan := testAWSSQSBroker.GetStopReceivingChanForTest()
		firstStep <- 1
		stopReceivingChan <- 1
	}()

	var (
		whetherContinue bool
		err             error
	)
	<-firstStep
	// Test the case that a signal was received from stopReceivingChan
	go func() {
		whetherContinue, err = testAWSSQSBroker.ContinueReceivingMessagesForTest(qURL, deliveries)
		nextStep <- 1
	}()
	<-nextStep
	assert.False(t, whetherContinue)
	assert.Nil(t, err)

	// Test the default condition
	whetherContinue, err = testAWSSQSBroker.ContinueReceivingMessagesForTest(qURL, deliveries)
	assert.True(t, whetherContinue)
	assert.Nil(t, err)

	// Test the error
	whetherContinue, err = errAWSSQSBroker.ContinueReceivingMessagesForTest(qURL, deliveries)
	assert.True(t, whetherContinue)
	assert.NotNil(t, err)

	// Test when there is no message
	outputCopy := *receiveMessageOutput
	receiveMessageOutput.Messages = []*sqs.Message{}
	whetherContinue, err = testAWSSQSBroker.ContinueReceivingMessagesForTest(qURL, deliveries)
	assert.True(t, whetherContinue)
	assert.Nil(t, err)
	// recover original value
	*receiveMessageOutput = outputCopy

}

func TestPrivateFunc_consume(t *testing.T) {
	server1, err := machinery.NewServer(cnf)
	if err != nil {
		t.Fatal(err)
	}
	wk := server1.NewWorker("sms_worker", 0)
	deliveries := make(chan *sqs.ReceiveMessageOutput)
	outputCopy := *receiveMessageOutput
	outputCopy.Messages = []*sqs.Message{}
	go func() { deliveries <- &outputCopy }()

	// an infinite loop will be executed only when there is no error
	err = testAWSSQSBroker.ConsumeForTest(deliveries, 0, wk)
	assert.NotNil(t, err)

}

func TestPrivateFunc_consumeOne(t *testing.T) {
	server1, err := machinery.NewServer(cnf)
	if err != nil {
		t.Fatal(err)
	}
	wk := server1.NewWorker("sms_worker", 0)
	err = testAWSSQSBroker.ConsumeOneForTest(receiveMessageOutput, wk)
	assert.NotNil(t, err)

	outputCopy := *receiveMessageOutput
	outputCopy.Messages = []*sqs.Message{}
	err = testAWSSQSBroker.ConsumeOneForTest(&outputCopy, wk)
	assert.NotNil(t, err)

	outputCopy.Messages = []*sqs.Message{
		{
			Body: aws.String("foo message"),
		},
	}
	err = testAWSSQSBroker.ConsumeOneForTest(&outputCopy, wk)
	assert.NotNil(t, err)
}

func TestPrivateFunc_initializePool(t *testing.T) {
	concurrency := 9
	pool := make(chan struct{}, concurrency)
	testAWSSQSBroker.InitializePoolForTest(pool, concurrency)
	assert.Len(t, pool, concurrency)
}

func TestPrivateFunc_startConsuming(t *testing.T) {
	server1, err := machinery.NewServer(cnf)
	if err != nil {
		t.Fatal(err)
	}

	wk := server1.NewWorker("sms_worker", 0)
	retryFunc := testAWSSQSBroker.GetRetryFuncForTest()
	stopChan := testAWSSQSBroker.GetStopChanForTest()
	retryStopChan := testAWSSQSBroker.GetRetryStopChanForTest()
	assert.Nil(t, retryFunc)
	testAWSSQSBroker.StartConsumingForTest("fooTag", wk)
	assert.IsType(t, retryFunc, retry.Closure())
	assert.Equal(t, len(stopChan), 0)
	assert.Equal(t, len(retryStopChan), 0)
}

func TestPrivateFuncDefaultQueueURL(t *testing.T) {
	qURL := testAWSSQSBroker.DefaultQueueURLForTest()

	assert.EqualValues(t, *qURL, "https://sqs.foo.amazonaws.com.cn/test_queue")
}

func TestPrivateFunc_stopReceiving(t *testing.T) {
	go testAWSSQSBroker.StopReceivingForTest()
	stopReceivingChan := testAWSSQSBroker.GetStopReceivingChanForTest()
	assert.NotNil(t, <-stopReceivingChan)
}

func TestPrivateFunc_receiveMessage(t *testing.T) {
	qURL := testAWSSQSBroker.DefaultQueueURLForTest()
	output, err := testAWSSQSBroker.ReceiveMessageForTest(qURL)
	assert.Nil(t, err)
	assert.Equal(t, receiveMessageOutput, output)
}

func TestPrivateFunc_consumeDeliveries(t *testing.T) {
	concurrency := 0
	pool := make(chan struct{}, concurrency)
	errorsChan := make(chan error)
	deliveries := make(chan *sqs.ReceiveMessageOutput, 0)
	server1, err := machinery.NewServer(cnf)
	if err != nil {
		t.Fatal(err)
	}
	wk := server1.NewWorker("sms_worker", 0)
	go func() { deliveries <- receiveMessageOutput }()
	whetherContinue, err := testAWSSQSBroker.ConsumeDeliveriesForTest(deliveries, concurrency, wk, pool, errorsChan)
	assert.True(t, whetherContinue)
	assert.Nil(t, err)

	go func() { errorsChan <- errors.New("foo error") }()
	whetherContinue, err = testAWSSQSBroker.ConsumeDeliveriesForTest(deliveries, concurrency, wk, pool, errorsChan)
	assert.False(t, whetherContinue)
	assert.NotNil(t, err)

	go func() { testAWSSQSBroker.GetStopChanForTest() <- 1 }()
	whetherContinue, err = testAWSSQSBroker.ConsumeDeliveriesForTest(deliveries, concurrency, wk, pool, errorsChan)
	assert.False(t, whetherContinue)
	assert.Nil(t, err)

	outputCopy := *receiveMessageOutput
	outputCopy.Messages = []*sqs.Message{}
	go func() { deliveries <- &outputCopy }()
	whetherContinue, err = testAWSSQSBroker.ConsumeDeliveriesForTest(deliveries, concurrency, wk, pool, errorsChan)
	e := <-errorsChan
	assert.True(t, whetherContinue)
	assert.NotNil(t, e)
	assert.Nil(t, err)

	go func() { deliveries <- receiveMessageOutput }()
	// wait for a while to fix racing problem
	time.Sleep(5 * time.Second)
	go func() { pool <- struct{}{} }()
	whetherContinue, err = testAWSSQSBroker.ConsumeDeliveriesForTest(deliveries, concurrency, wk, pool, errorsChan)
	p := <-pool
	assert.True(t, whetherContinue)
	assert.NotNil(t, p)
	assert.Nil(t, err)
}

func TestPrivateFunc_deleteOne(t *testing.T) {
	err := testAWSSQSBroker.DeleteOneForTest(receiveMessageOutput)
	assert.Nil(t, err)

	err = errAWSSQSBroker.DeleteOneForTest(receiveMessageOutput)
	assert.NotNil(t, err)
}
<commit_msg>rm channel's capacity<commit_after>package brokers_test

import (
	"errors"
	"github.com/RichardKnop/machinery/v1"
	"github.com/RichardKnop/machinery/v1/brokers"
	"github.com/RichardKnop/machinery/v1/config"
	"github.com/RichardKnop/machinery/v1/retry"
	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/service/sqs"
	"github.com/stretchr/testify/assert"
	"testing"
	"time"
)

var (
	testBroker           brokers.Interface
	testAWSSQSBroker     *brokers.AWSSQSBroker
	errAWSSQSBroker      *brokers.AWSSQSBroker
	cnf                  *config.Config
	receiveMessageOutput *sqs.ReceiveMessageOutput
)

func init() {
	testAWSSQSBroker = brokers.TestAWSSQSBroker
	errAWSSQSBroker = brokers.ErrAWSSQSBroker
	cnf = brokers.TestConf
	receiveMessageOutput = brokers.ReceiveMessageOutput
	testBroker = brokers.NewAWSSQSBroker(cnf)
}

func TestNewAWSSQSBroker(t *testing.T) {
	assert.IsType(t, testAWSSQSBroker, testBroker)
}

func TestPrivateFunc_continueReceivingMessages(t *testing.T) {
	qURL := testAWSSQSBroker.DefaultQueueURLForTest()
	deliveries := make(chan *sqs.ReceiveMessageOutput)
	firstStep := make(chan int)
	nextStep := make(chan int)
	go func() {
		stopReceivingChan := testAWSSQSBroker.GetStopReceivingChanForTest()
		firstStep <- 1
		stopReceivingChan <- 1
	}()

	var (
		whetherContinue bool
		err             error
	)
	<-firstStep
	// Test the case that a signal was received from stopReceivingChan
	go func() {
		whetherContinue, err = testAWSSQSBroker.ContinueReceivingMessagesForTest(qURL, deliveries)
		nextStep <- 1
	}()
	<-nextStep
	assert.False(t, whetherContinue)
	assert.Nil(t, err)

	// Test the default condition
	whetherContinue, err = testAWSSQSBroker.ContinueReceivingMessagesForTest(qURL, deliveries)
	assert.True(t, whetherContinue)
	assert.Nil(t, err)

	// Test the error
	whetherContinue, err = errAWSSQSBroker.ContinueReceivingMessagesForTest(qURL, deliveries)
	assert.True(t, whetherContinue)
	assert.NotNil(t, err)

	// Test when there is no message
	outputCopy := *receiveMessageOutput
	receiveMessageOutput.Messages = []*sqs.Message{}
	whetherContinue, err = testAWSSQSBroker.ContinueReceivingMessagesForTest(qURL, deliveries)
	assert.True(t, whetherContinue)
	assert.Nil(t, err)
	// recover original value
	*receiveMessageOutput = outputCopy

}

func TestPrivateFunc_consume(t *testing.T) {
	server1, err := machinery.NewServer(cnf)
	if err != nil {
		t.Fatal(err)
	}
	wk := server1.NewWorker("sms_worker", 0)
	deliveries := make(chan *sqs.ReceiveMessageOutput)
	outputCopy := *receiveMessageOutput
	outputCopy.Messages = []*sqs.Message{}
	go func() { deliveries <- &outputCopy }()

	// an infinite loop will be executed only when there is no error
	err = testAWSSQSBroker.ConsumeForTest(deliveries, 0, wk)
	assert.NotNil(t, err)

}

func TestPrivateFunc_consumeOne(t *testing.T) {
	server1, err := machinery.NewServer(cnf)
	if err != nil {
		t.Fatal(err)
	}
	wk := server1.NewWorker("sms_worker", 0)
	err = testAWSSQSBroker.ConsumeOneForTest(receiveMessageOutput, wk)
	assert.NotNil(t, err)

	outputCopy := *receiveMessageOutput
	outputCopy.Messages = []*sqs.Message{}
	err = testAWSSQSBroker.ConsumeOneForTest(&outputCopy, wk)
	assert.NotNil(t, err)

	outputCopy.Messages = []*sqs.Message{
		{
			Body: aws.String("foo message"),
		},
	}
	err = testAWSSQSBroker.ConsumeOneForTest(&outputCopy, wk)
	assert.NotNil(t, err)
}

func TestPrivateFunc_initializePool(t *testing.T) {
	concurrency := 9
	pool := make(chan struct{}, concurrency)
	testAWSSQSBroker.InitializePoolForTest(pool, concurrency)
	assert.Len(t, pool, concurrency)
}

func TestPrivateFunc_startConsuming(t *testing.T) {
	server1, err := machinery.NewServer(cnf)
	if err != nil {
		t.Fatal(err)
	}

	wk := server1.NewWorker("sms_worker", 0)
	retryFunc := testAWSSQSBroker.GetRetryFuncForTest()
	stopChan := testAWSSQSBroker.GetStopChanForTest()
	retryStopChan := testAWSSQSBroker.GetRetryStopChanForTest()
	assert.Nil(t, retryFunc)
	testAWSSQSBroker.StartConsumingForTest("fooTag", wk)
	assert.IsType(t, retryFunc, retry.Closure())
	assert.Equal(t, len(stopChan), 0)
	assert.Equal(t, len(retryStopChan), 0)
}

func TestPrivateFuncDefaultQueueURL(t *testing.T) {
	qURL := testAWSSQSBroker.DefaultQueueURLForTest()

	assert.EqualValues(t, *qURL, "https://sqs.foo.amazonaws.com.cn/test_queue")
}

func TestPrivateFunc_stopReceiving(t *testing.T) {
	go testAWSSQSBroker.StopReceivingForTest()
	stopReceivingChan := testAWSSQSBroker.GetStopReceivingChanForTest()
	assert.NotNil(t, <-stopReceivingChan)
}

func TestPrivateFunc_receiveMessage(t *testing.T) {
	qURL := testAWSSQSBroker.DefaultQueueURLForTest()
	output, err := testAWSSQSBroker.ReceiveMessageForTest(qURL)
	assert.Nil(t, err)
	assert.Equal(t, receiveMessageOutput, output)
}

func TestPrivateFunc_consumeDeliveries(t *testing.T) {
	concurrency := 0
	pool := make(chan struct{}, concurrency)
	errorsChan := make(chan error)
	deliveries := make(chan *sqs.ReceiveMessageOutput)
	server1, err := machinery.NewServer(cnf)
	if err != nil {
		t.Fatal(err)
	}
	wk := server1.NewWorker("sms_worker", 0)
	go func() { deliveries <- receiveMessageOutput }()
	whetherContinue, err := testAWSSQSBroker.ConsumeDeliveriesForTest(deliveries, concurrency, wk, pool, errorsChan)
	assert.True(t, whetherContinue)
	assert.Nil(t, err)

	go func() { errorsChan <- errors.New("foo error") }()
	whetherContinue, err = testAWSSQSBroker.ConsumeDeliveriesForTest(deliveries, concurrency, wk, pool, errorsChan)
	assert.False(t, whetherContinue)
	assert.NotNil(t, err)

	go func() { testAWSSQSBroker.GetStopChanForTest() <- 1 }()
	whetherContinue, err = testAWSSQSBroker.ConsumeDeliveriesForTest(deliveries, concurrency, wk, pool, errorsChan)
	assert.False(t, whetherContinue)
	assert.Nil(t, err)

	outputCopy := *receiveMessageOutput
	outputCopy.Messages = []*sqs.Message{}
	go func() { deliveries <- &outputCopy }()
	whetherContinue, err = testAWSSQSBroker.ConsumeDeliveriesForTest(deliveries, concurrency, wk, pool, errorsChan)
	e := <-errorsChan
	assert.True(t, whetherContinue)
	assert.NotNil(t, e)
	assert.Nil(t, err)

	go func() { deliveries <- receiveMessageOutput }()
	// wait for a while to fix racing problem
	time.Sleep(5 * time.Second)
	go func() { pool <- struct{}{} }()
	whetherContinue, err = testAWSSQSBroker.ConsumeDeliveriesForTest(deliveries, concurrency, wk, pool, errorsChan)
	p := <-pool
	assert.True(t, whetherContinue)
	assert.NotNil(t, p)
	assert.Nil(t, err)
}

func TestPrivateFunc_deleteOne(t *testing.T) {
	err := testAWSSQSBroker.DeleteOneForTest(receiveMessageOutput)
	assert.Nil(t, err)

	err = errAWSSQSBroker.DeleteOneForTest(receiveMessageOutput)
	assert.NotNil(t, err)
}
<|endoftext|>


<commit_before>//
// mod_http.go
//
// http mod panel
//

package srnd

import (
  "github.com/majestrate/srndv2/src/nacl"
  "github.com/gorilla/sessions"
  "errors"
  "encoding/hex"
  "encoding/json"
  "fmt"
  "io"
  "log"
  "net/http"
  "strings"
)

type httpModUI struct {
  regenAll func()
  regen func (ArticleEntry)
  delete func (string)
  modMessageChan chan NNTPMessage
  database Database
  articles ArticleStore
  store *sessions.CookieStore
  prefix string
  mod_prefix string
}

func createHttpModUI(frontend httpFrontend) httpModUI {
  return httpModUI{frontend.regenAll, frontend.Regen, frontend.deleteThreadMarkup, make(chan NNTPMessage), frontend.daemon.database, frontend.daemon.store, frontend.store, frontend.prefix, frontend.prefix + "mod/"}

}
func (self httpModUI) getAdminFunc(funcname string) AdminFunc {
  if funcname == "template.reload" {
    return func(param map[string]interface{}) (string, error) {
      tname, ok := param["template"]
      if ok {
        t := ""
        switch tname.(type) {
        case string:
          t = tname.(string)
        default:
          return "failed to reload templates", errors.New("invalid parameters")
        }
        template.reloadTemplate(t)
        return "reloaded " + t, nil
      }
      template.reloadAllTemplates()
      return "reloaded all templates", nil
    }
  } else if funcname == "frontend.regen" {
    return func(param map[string]interface{}) (string, error) {
      go self.regenAll()
      return "started regeneration", nil
    }
  } else if funcname == "thumbnail.regen" {
    return func(param map[string]interface{}) (string, error) {
      threads, ok := param["threads"]
      t := 4
      if ok {
        switch threads.(type) {
        case int64:
          t = int(threads.(int64))
          if t <= 0 {
            return "failed to regen thumbnails", errors.New("invalid number of threads")
          }
        default:
          return "failed to regen thumbnails", errors.New("invalid parameters")
        }
      }
      go reThumbnail(t, self.articles)
      return fmt.Sprintf("started rethumbnailing with %d threads", t), nil
    }
  }
  return nil
}

// handle an admin action
func (self httpModUI) HandleAdminCommand(wr http.ResponseWriter, r *http.Request) {
  self.asAuthed(func(url string) {
    action := strings.Split(url, "/admin/")[1]
    f := self.getAdminFunc(action)
    if f == nil {
      wr.WriteHeader(404)
    } else {
      var msg string
      var err error
      req := make(map[string]interface{})
      if r.Method == "POST" {
        dec := json.NewDecoder(r.Body)
        err = dec.Decode(req)
      }
      if err == nil {
        msg, err = f(req)
      }
      resp := make(map[string]interface{})
      if err == nil {
        resp["error"] = nil
      } else {
        resp["error"] = err.Error()
      }
      resp["result"] = msg
      enc := json.NewEncoder(wr)
      enc.Encode(resp)
    }

  }, wr, r)
}

// TODO: check for different levels of permissions
func (self httpModUI) CheckKey(privkey string) (bool, error) {
  privkey_bytes, err := hex.DecodeString(privkey)
  if err == nil {
    pubkey_bytes := nacl.GetSignPubkey(privkey_bytes)
    if pubkey_bytes != nil {
      pubkey := hex.EncodeToString(pubkey_bytes)
      if self.database.CheckModPubkeyGlobal(pubkey) {
        // this user is an admin
        return true, nil
      } else {
        return false, nil
      }
    }
  }
  log.Println("invalid key format for key", privkey)
  return false, err
}

func (self httpModUI) MessageChan() chan NNTPMessage {
  return self.modMessageChan
}

func (self httpModUI) getSession(r *http.Request) *sessions.Session {
  s, _ := self.store.Get(r, "nntpchan-mod")
  return s
}

// get the session's private key as bytes or nil if we don't have it
func (self httpModUI) getSessionPrivkeyBytes(r *http.Request) []byte {
  s := self.getSession(r)
  k, ok := s.Values["privkey"]
  if ok {
    privkey_bytes, err := hex.DecodeString(k.(string))
    if err == nil {
      return privkey_bytes
    }
    log.Println("failed to decode private key bytes from session", err)
  } else {
    log.Println("failed to get private key from session, no private key in session?")
  }
  return nil
}

// returns true if the session is okay
// otherwise redirect to login page
func (self httpModUI) checkSession(r *http.Request) bool {
  s := self.getSession(r)
  k, ok := s.Values["privkey"]
  if ok {
    ok, err := self.CheckKey(k.(string))
    if err != nil {
      return false
    }
    return ok
  }
  return false
}

func (self httpModUI) writeTemplate(wr http.ResponseWriter, name string) {
  self.writeTemplateParam(wr, name, nil)
}

func (self httpModUI) writeTemplateParam(wr http.ResponseWriter, name string, param map[string]string) {
  if param == nil {
    param = make(map[string]string)
  }
  param["prefix"] = self.prefix
  param["mod_prefix"] = self.mod_prefix
  io.WriteString(wr, template.renderTemplate(name, param))
}

// do a function as authenticated
// pass in the request path to the handler
func (self httpModUI) asAuthed(handler func(string), wr http.ResponseWriter, r *http.Request) {
  if self.checkSession(r) {
    handler(r.URL.Path)
  } else {
    wr.WriteHeader(403)
  }
}

// do stuff to a certain message if with have it and are authed
func (self httpModUI) asAuthedWithMessage(handler func(ArticleEntry, *http.Request) map[string]interface{}, wr http.ResponseWriter, req *http.Request) {
  self.asAuthed(func(path string) {
    // get the long hash
    if strings.Count(path, "/") > 2 {
      // TOOD: prefix detection
      longhash := strings.Split(path, "/")[3]
      // get the message id
      msg, err := self.database.GetMessageIDByHash(longhash)
      resp := make(map[string]interface{})
      if err == nil {
        resp = handler(msg, req)
      } else {
        resp["error"] = fmt.Sprintf("don't have message with hash %s, %s", longhash, err.Error())
      }
      enc := json.NewEncoder(wr)
      enc.Encode(resp)
    } else {
      wr.WriteHeader(404)
    }
  }, wr, req)
}

func (self httpModUI) HandleAddPubkey(wr http.ResponseWriter, r *http.Request) {
}

func (self httpModUI) HandleDelPubkey(wr http.ResponseWriter, r *http.Request) {
}

func (self httpModUI) HandleUnbanAddress(wr http.ResponseWriter, r *http.Request) {
  self.asAuthed(func(path string) {
    // extract the ip address
    // TODO: ip ranges and prefix detection
    if strings.Count(path, "/") > 2 {
      addr := strings.Split(path, "/")[3]
      resp := make(map[string]interface{})
      banned, err := self.database.CheckIPBanned(addr)
      if err != nil {
        resp["error"] = fmt.Sprintf("cannot tell if %s is banned: %s", addr, err.Error())
      } else if banned {
        // TODO: rangebans
        err = self.database.UnbanAddr(addr)
        if err == nil {
          resp["result"] = fmt.Sprintf("%s was unbanned", addr)
        } else {
          resp["error"] = err.Error()
        }
      } else {
        resp["error"] = fmt.Sprintf("%s was not banned", addr)
      }
      enc := json.NewEncoder(wr)
      enc.Encode(resp)
    } else {
      wr.WriteHeader(404)
    }
  }, wr, r)
}

// handle ban logic
func (self httpModUI) handleBanAddress(msg ArticleEntry, r *http.Request) map[string]interface{} {
  // get the article headers
  resp := make(map[string]interface{})
  msgid := msg.MessageID()
  hdr := self.articles.GetHeaders(msgid)
  if hdr == nil {
    // we don't got it?!
    resp["error"] = fmt.Sprintf("message %s not on the filesystem wtf?", msgid)
  } else {
    // get the associated encrypted ip
    encip := hdr.Get("X-Encrypted-Ip", hdr.Get("X-Encrypted-IP", ""))
    encip = strings.Trim(encip, "\t ")

    if len(encip) == 0 {
      // no ip header detected
      resp["error"] = fmt.Sprintf("%s has no IP, ban Tor instead", msgid)
    } else {
      // get the ip address if we have it
      ip, err := self.database.GetIPAddress(encip)
      if len(ip) > 0 {
        // we have it
        // ban the address
        err = self.database.BanAddr(ip)
        // then we tell everyone about it
        var key string
        // TODO: we SHOULD have the key, but what if we do not?
        key, err = self.database.GetEncKey(encip)
        // create mod message
        // TODO: hardcoded ban period
        mm := ModMessage{overchanInetBan(encip, key, -1),}
        privkey_bytes := self.getSessionPrivkeyBytes(r)
        if privkey_bytes == nil {
          // this should not happen
          log.Println("failed to get privkey bytes from session")
          resp["error"] = "failed to get private key from session. wtf?"
        } else {
          // wrap and sign
          nntp := wrapModMessage(mm)
          nntp, err = signArticle(nntp, privkey_bytes)
          if err == nil {
            // federate
            self.modMessageChan <- nntp
          }
        }
      } else {
        // we don't have it
        // ban the encrypted version
        err = self.database.BanEncAddr(encip)
      }
      if err == nil {
        result_msg :=  fmt.Sprintf("We banned %s", encip)
        if len(ip) > 0 {
          result_msg += fmt.Sprintf(" (%s)", ip)
        }
        resp["banned"] = result_msg
      } else {
        resp["error"] = err.Error()
      }

    }
  }
  return resp
}

func (self httpModUI) handleDeletePost(msg ArticleEntry, r *http.Request) map[string]interface{} {
  var mm ModMessage
  resp := make(map[string]interface{})
  msgid := msg.MessageID()
  delmsgs := []string{}
  // get headers
  hdr := self.articles.GetHeaders(msgid)
  if hdr == nil {
    resp["error"] = fmt.Sprintf("message %s is not on the filesystem? wtf!", msgid)
  } else {
    ref := hdr.Get("References", hdr.Get("Reference", ""))
    ref = strings.Trim(ref, "\t ")
    // is it a root post?
    if ref == "" {
      // load replies
      replies := self.database.GetThreadReplies(msgid, 0)
      if replies != nil {
        for _, repl := range(replies) {
          // append mod line to mod message for reply
          mm = append(mm, overchanDelete(repl))
          // add to delete queue
          delmsgs = append(delmsgs, repl)
        }
      }
    }
    delmsgs = append(delmsgs, msgid)
    // append mod line to mod message
    mm = append(mm, overchanDelete(msgid))

    resp["deleted"] = delmsgs
    // only regen threads when we delete a non root port
    if ref != "" {
      group := hdr.Get("Newsgroups", "")
      self.regen(ArticleEntry{
        ref, group,
      })
    }
    privkey_bytes := self.getSessionPrivkeyBytes(r)
    if privkey_bytes == nil {
      // crap this should never happen
      log.Println("failed to get private keys from session, not federating")
    } else {
      // wrap and sign mod message
      nntp, err := signArticle(wrapModMessage(mm), privkey_bytes)
      if err == nil {
        // send it off to federate
        self.modMessageChan <- nntp
      } else {
        resp["error"] = fmt.Sprintf("signing error: %s", err.Error())
      }
    }
  }
  return resp
}

// ban the address of a poster
func (self httpModUI) HandleBanAddress(wr http.ResponseWriter, r *http.Request) {
  self.asAuthedWithMessage(self.handleBanAddress, wr, r)
}

// delete a post
func (self httpModUI) HandleDeletePost(wr http.ResponseWriter, r *http.Request) {
  self.asAuthedWithMessage(self.handleDeletePost, wr, r)
}

func (self httpModUI) HandleLogin(wr http.ResponseWriter, r *http.Request) {
  privkey := r.FormValue("privkey")
  msg := "failed login: "
  if len(privkey) == 0 {
    msg += "no key"
  } else {
    ok, err := self.CheckKey(privkey)
    if err != nil {
      msg += fmt.Sprintf("%s", err)
    } else if ok {
      msg = "login okay"
      sess := self.getSession(r)
      sess.Values["privkey"] = privkey
      sess.Save(r, wr)
    } else {
      msg += "invalid key"
    }
  }
  self.writeTemplateParam(wr, "modlogin_result.mustache", map[string]string { "message" : msg })
}

func (self httpModUI) HandleKeyGen(wr http.ResponseWriter, r *http.Request) {
  pk, sk := newSignKeypair()
  tripcode := makeTripcode(pk)
  self.writeTemplateParam(wr, "keygen.mustache", map[string]string {"public" : pk, "secret" : sk, "tripcode" : tripcode})
}

func (self httpModUI) ServeModPage(wr http.ResponseWriter, r *http.Request) {
  if self.checkSession(r) {
    // we are logged in
    // serve mod page
    self.writeTemplate(wr, "modpage.mustache")
  } else {
    // we are not logged in
    // serve login page
    self.writeTemplate(wr, "modlogin.mustache")
  }

}
<commit_msg>add logging<commit_after>//
// mod_http.go
//
// http mod panel
//

package srnd

import (
  "github.com/majestrate/srndv2/src/nacl"
  "github.com/gorilla/sessions"
  "errors"
  "encoding/hex"
  "encoding/json"
  "fmt"
  "io"
  "log"
  "net/http"
  "strings"
)

type httpModUI struct {
  regenAll func()
  regen func (ArticleEntry)
  delete func (string)
  modMessageChan chan NNTPMessage
  database Database
  articles ArticleStore
  store *sessions.CookieStore
  prefix string
  mod_prefix string
}

func createHttpModUI(frontend httpFrontend) httpModUI {
  return httpModUI{frontend.regenAll, frontend.Regen, frontend.deleteThreadMarkup, make(chan NNTPMessage), frontend.daemon.database, frontend.daemon.store, frontend.store, frontend.prefix, frontend.prefix + "mod/"}

}
func (self httpModUI) getAdminFunc(funcname string) AdminFunc {
  if funcname == "template.reload" {
    return func(param map[string]interface{}) (string, error) {
      tname, ok := param["template"]
      if ok {
        t := ""
        switch tname.(type) {
        case string:
          t = tname.(string)
        default:
          return "failed to reload templates", errors.New("invalid parameters")
        }
        template.reloadTemplate(t)
        return "reloaded " + t, nil
      }
      template.reloadAllTemplates()
      return "reloaded all templates", nil
    }
  } else if funcname == "frontend.regen" {
    return func(param map[string]interface{}) (string, error) {
      go self.regenAll()
      return "started regeneration", nil
    }
  } else if funcname == "thumbnail.regen" {
    return func(param map[string]interface{}) (string, error) {
      threads, ok := param["threads"]
      t := 4
      if ok {
        switch threads.(type) {
        case int64:
          t = int(threads.(int64))
          if t <= 0 {
            return "failed to regen thumbnails", errors.New("invalid number of threads")
          }
        default:
          return "failed to regen thumbnails", errors.New("invalid parameters")
        }
      }
      go reThumbnail(t, self.articles)
      return fmt.Sprintf("started rethumbnailing with %d threads", t), nil
    }
  }
  return nil
}

// handle an admin action
func (self httpModUI) HandleAdminCommand(wr http.ResponseWriter, r *http.Request) {
  self.asAuthed(func(url string) {
    action := strings.Split(url, "/admin/")[1]
    log.Println("try admin action", action)
    f := self.getAdminFunc(action)
    if f == nil {
      wr.WriteHeader(404)
    } else {
      var msg string
      var err error
      req := make(map[string]interface{})
      if r.Method == "POST" {
        dec := json.NewDecoder(r.Body)
        err = dec.Decode(req)
      }
      if err == nil {
        msg, err = f(req)
      }
      resp := make(map[string]interface{})
      if err == nil {
        resp["error"] = nil
      } else {
        resp["error"] = err.Error()
      }
      resp["result"] = msg
      enc := json.NewEncoder(wr)
      enc.Encode(resp)
    }

  }, wr, r)
}

// TODO: check for different levels of permissions
func (self httpModUI) CheckKey(privkey string) (bool, error) {
  privkey_bytes, err := hex.DecodeString(privkey)
  if err == nil {
    pubkey_bytes := nacl.GetSignPubkey(privkey_bytes)
    if pubkey_bytes != nil {
      pubkey := hex.EncodeToString(pubkey_bytes)
      if self.database.CheckModPubkeyGlobal(pubkey) {
        // this user is an admin
        return true, nil
      } else {
        return false, nil
      }
    }
  }
  log.Println("invalid key format for key", privkey)
  return false, err
}

func (self httpModUI) MessageChan() chan NNTPMessage {
  return self.modMessageChan
}

func (self httpModUI) getSession(r *http.Request) *sessions.Session {
  s, _ := self.store.Get(r, "nntpchan-mod")
  return s
}

// get the session's private key as bytes or nil if we don't have it
func (self httpModUI) getSessionPrivkeyBytes(r *http.Request) []byte {
  s := self.getSession(r)
  k, ok := s.Values["privkey"]
  if ok {
    privkey_bytes, err := hex.DecodeString(k.(string))
    if err == nil {
      return privkey_bytes
    }
    log.Println("failed to decode private key bytes from session", err)
  } else {
    log.Println("failed to get private key from session, no private key in session?")
  }
  return nil
}

// returns true if the session is okay
// otherwise redirect to login page
func (self httpModUI) checkSession(r *http.Request) bool {
  s := self.getSession(r)
  k, ok := s.Values["privkey"]
  if ok {
    ok, err := self.CheckKey(k.(string))
    if err != nil {
      return false
    }
    return ok
  }
  return false
}

func (self httpModUI) writeTemplate(wr http.ResponseWriter, name string) {
  self.writeTemplateParam(wr, name, nil)
}

func (self httpModUI) writeTemplateParam(wr http.ResponseWriter, name string, param map[string]string) {
  if param == nil {
    param = make(map[string]string)
  }
  param["prefix"] = self.prefix
  param["mod_prefix"] = self.mod_prefix
  io.WriteString(wr, template.renderTemplate(name, param))
}

// do a function as authenticated
// pass in the request path to the handler
func (self httpModUI) asAuthed(handler func(string), wr http.ResponseWriter, r *http.Request) {
  if self.checkSession(r) {
    handler(r.URL.Path)
  } else {
    wr.WriteHeader(403)
  }
}

// do stuff to a certain message if with have it and are authed
func (self httpModUI) asAuthedWithMessage(handler func(ArticleEntry, *http.Request) map[string]interface{}, wr http.ResponseWriter, req *http.Request) {
  self.asAuthed(func(path string) {
    // get the long hash
    if strings.Count(path, "/") > 2 {
      // TOOD: prefix detection
      longhash := strings.Split(path, "/")[3]
      // get the message id
      msg, err := self.database.GetMessageIDByHash(longhash)
      resp := make(map[string]interface{})
      if err == nil {
        resp = handler(msg, req)
      } else {
        resp["error"] = fmt.Sprintf("don't have message with hash %s, %s", longhash, err.Error())
      }
      enc := json.NewEncoder(wr)
      enc.Encode(resp)
    } else {
      wr.WriteHeader(404)
    }
  }, wr, req)
}

func (self httpModUI) HandleAddPubkey(wr http.ResponseWriter, r *http.Request) {
}

func (self httpModUI) HandleDelPubkey(wr http.ResponseWriter, r *http.Request) {
}

func (self httpModUI) HandleUnbanAddress(wr http.ResponseWriter, r *http.Request) {
  self.asAuthed(func(path string) {
    // extract the ip address
    // TODO: ip ranges and prefix detection
    if strings.Count(path, "/") > 2 {
      addr := strings.Split(path, "/")[3]
      resp := make(map[string]interface{})
      banned, err := self.database.CheckIPBanned(addr)
      if err != nil {
        resp["error"] = fmt.Sprintf("cannot tell if %s is banned: %s", addr, err.Error())
      } else if banned {
        // TODO: rangebans
        err = self.database.UnbanAddr(addr)
        if err == nil {
          resp["result"] = fmt.Sprintf("%s was unbanned", addr)
        } else {
          resp["error"] = err.Error()
        }
      } else {
        resp["error"] = fmt.Sprintf("%s was not banned", addr)
      }
      enc := json.NewEncoder(wr)
      enc.Encode(resp)
    } else {
      wr.WriteHeader(404)
    }
  }, wr, r)
}

// handle ban logic
func (self httpModUI) handleBanAddress(msg ArticleEntry, r *http.Request) map[string]interface{} {
  // get the article headers
  resp := make(map[string]interface{})
  msgid := msg.MessageID()
  hdr := self.articles.GetHeaders(msgid)
  if hdr == nil {
    // we don't got it?!
    resp["error"] = fmt.Sprintf("message %s not on the filesystem wtf?", msgid)
  } else {
    // get the associated encrypted ip
    encip := hdr.Get("X-Encrypted-Ip", hdr.Get("X-Encrypted-IP", ""))
    encip = strings.Trim(encip, "\t ")

    if len(encip) == 0 {
      // no ip header detected
      resp["error"] = fmt.Sprintf("%s has no IP, ban Tor instead", msgid)
    } else {
      // get the ip address if we have it
      ip, err := self.database.GetIPAddress(encip)
      if len(ip) > 0 {
        // we have it
        // ban the address
        err = self.database.BanAddr(ip)
        // then we tell everyone about it
        var key string
        // TODO: we SHOULD have the key, but what if we do not?
        key, err = self.database.GetEncKey(encip)
        // create mod message
        // TODO: hardcoded ban period
        mm := ModMessage{overchanInetBan(encip, key, -1),}
        privkey_bytes := self.getSessionPrivkeyBytes(r)
        if privkey_bytes == nil {
          // this should not happen
          log.Println("failed to get privkey bytes from session")
          resp["error"] = "failed to get private key from session. wtf?"
        } else {
          // wrap and sign
          nntp := wrapModMessage(mm)
          nntp, err = signArticle(nntp, privkey_bytes)
          if err == nil {
            // federate
            self.modMessageChan <- nntp
          }
        }
      } else {
        // we don't have it
        // ban the encrypted version
        err = self.database.BanEncAddr(encip)
      }
      if err == nil {
        result_msg :=  fmt.Sprintf("We banned %s", encip)
        if len(ip) > 0 {
          result_msg += fmt.Sprintf(" (%s)", ip)
        }
        resp["banned"] = result_msg
      } else {
        resp["error"] = err.Error()
      }

    }
  }
  return resp
}

func (self httpModUI) handleDeletePost(msg ArticleEntry, r *http.Request) map[string]interface{} {
  var mm ModMessage
  resp := make(map[string]interface{})
  msgid := msg.MessageID()
  delmsgs := []string{}
  // get headers
  hdr := self.articles.GetHeaders(msgid)
  if hdr == nil {
    resp["error"] = fmt.Sprintf("message %s is not on the filesystem? wtf!", msgid)
  } else {
    ref := hdr.Get("References", hdr.Get("Reference", ""))
    ref = strings.Trim(ref, "\t ")
    // is it a root post?
    if ref == "" {
      // load replies
      replies := self.database.GetThreadReplies(msgid, 0)
      if replies != nil {
        for _, repl := range(replies) {
          // append mod line to mod message for reply
          mm = append(mm, overchanDelete(repl))
          // add to delete queue
          delmsgs = append(delmsgs, repl)
        }
      }
    }
    delmsgs = append(delmsgs, msgid)
    // append mod line to mod message
    mm = append(mm, overchanDelete(msgid))

    resp["deleted"] = delmsgs
    // only regen threads when we delete a non root port
    if ref != "" {
      group := hdr.Get("Newsgroups", "")
      self.regen(ArticleEntry{
        ref, group,
      })
    }
    privkey_bytes := self.getSessionPrivkeyBytes(r)
    if privkey_bytes == nil {
      // crap this should never happen
      log.Println("failed to get private keys from session, not federating")
    } else {
      // wrap and sign mod message
      nntp, err := signArticle(wrapModMessage(mm), privkey_bytes)
      if err == nil {
        // send it off to federate
        self.modMessageChan <- nntp
      } else {
        resp["error"] = fmt.Sprintf("signing error: %s", err.Error())
      }
    }
  }
  return resp
}

// ban the address of a poster
func (self httpModUI) HandleBanAddress(wr http.ResponseWriter, r *http.Request) {
  self.asAuthedWithMessage(self.handleBanAddress, wr, r)
}

// delete a post
func (self httpModUI) HandleDeletePost(wr http.ResponseWriter, r *http.Request) {
  self.asAuthedWithMessage(self.handleDeletePost, wr, r)
}

func (self httpModUI) HandleLogin(wr http.ResponseWriter, r *http.Request) {
  privkey := r.FormValue("privkey")
  msg := "failed login: "
  if len(privkey) == 0 {
    msg += "no key"
  } else {
    ok, err := self.CheckKey(privkey)
    if err != nil {
      msg += fmt.Sprintf("%s", err)
    } else if ok {
      msg = "login okay"
      sess := self.getSession(r)
      sess.Values["privkey"] = privkey
      sess.Save(r, wr)
    } else {
      msg += "invalid key"
    }
  }
  self.writeTemplateParam(wr, "modlogin_result.mustache", map[string]string { "message" : msg })
}

func (self httpModUI) HandleKeyGen(wr http.ResponseWriter, r *http.Request) {
  pk, sk := newSignKeypair()
  tripcode := makeTripcode(pk)
  self.writeTemplateParam(wr, "keygen.mustache", map[string]string {"public" : pk, "secret" : sk, "tripcode" : tripcode})
}

func (self httpModUI) ServeModPage(wr http.ResponseWriter, r *http.Request) {
  if self.checkSession(r) {
    // we are logged in
    // serve mod page
    self.writeTemplate(wr, "modpage.mustache")
  } else {
    // we are not logged in
    // serve login page
    self.writeTemplate(wr, "modlogin.mustache")
  }

}
<|endoftext|>


<commit_before>// Copyright (c) 2017 Intel Corporation
//
// 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.

// This is a "Multi-plugin".The delegate concept refered from CNI project
// It reads other plugin netconf, and then invoke them, e.g.
// flannel or sriov plugin.

package main

import (
	"encoding/json"
	"fmt"
	"io/ioutil"
	"os"
	"path/filepath"

	k8s "github.com/Intel-Corp/multus-cni/k8sclient"
	"github.com/Intel-Corp/multus-cni/types"
	"github.com/containernetworking/cni/pkg/invoke"
	"github.com/containernetworking/cni/pkg/skel"
	"github.com/containernetworking/cni/pkg/version"
)

const defaultCNIDir = "/var/lib/cni/multus"

var masterpluginEnabled bool
var defaultcninetwork bool

//taken from cni/plugins/meta/flannel/flannel.go
func isString(i interface{}) bool {
	_, ok := i.(string)
	return ok
}

func isBool(i interface{}) bool {
	_, ok := i.(bool)
	return ok
}

func loadNetConf(bytes []byte) (*types.NetConf, error) {
	netconf := &types.NetConf{}
	if err := json.Unmarshal(bytes, netconf); err != nil {
		return nil, fmt.Errorf("failed to load netconf: %v", err)
	}

	if netconf.Kubeconfig != "" && netconf.Delegates != nil {
		defaultcninetwork = true
	}

	if netconf.UseDefault {
		if netconf.Kubeconfig == "" || !defaultcninetwork {
			return nil, fmt.Errorf(`If you have set always_use_default, you must also set the delegates & the kubeconfig, refer to the README`)
		}
		return netconf, nil
	}

	if !netconf.UseDefault && (netconf.Kubeconfig != "" && !defaultcninetwork) {
		return netconf, nil
	}

	if len(netconf.Delegates) == 0 && !defaultcninetwork {
		return nil, fmt.Errorf(`delegates or kubeconfig option is must, refer README.md`)
	}

	if netconf.CNIDir == "" {
		netconf.CNIDir = defaultCNIDir
	}

	return netconf, nil
}

func saveScratchNetConf(containerID, dataDir string, netconf []byte) error {
	if err := os.MkdirAll(dataDir, 0700); err != nil {
		return fmt.Errorf("failed to create the multus data directory(%q): %v", dataDir, err)
	}

	path := filepath.Join(dataDir, containerID)

	err := ioutil.WriteFile(path, netconf, 0600)
	if err != nil {
		return fmt.Errorf("failed to write container data in the path(%q): %v", path, err)
	}

	return err
}

func consumeScratchNetConf(containerID, dataDir string) ([]byte, error) {
	path := filepath.Join(dataDir, containerID)
	defer os.Remove(path)

	return ioutil.ReadFile(path)
}

func getifname() (f func() string) {
	var interfaceIndex int
	f = func() string {
		ifname := fmt.Sprintf("net%d", interfaceIndex)
		interfaceIndex++
		return ifname
	}

	return
}

func saveDelegates(containerID, dataDir string, delegates []map[string]interface{}) error {
	delegatesBytes, err := json.Marshal(delegates)
	if err != nil {
		return fmt.Errorf("error serializing delegate netconf: %v", err)
	}

	if err = saveScratchNetConf(containerID, dataDir, delegatesBytes); err != nil {
		return fmt.Errorf("error in saving the  delegates : %v", err)
	}

	return err
}

func checkDelegate(netconf map[string]interface{}) error {
	if netconf["type"] == nil {
		return fmt.Errorf("delegate must have the field 'type'")
	}

	if !isString(netconf["type"]) {
		return fmt.Errorf("delegate field 'type' must be a string")
	}

	if netconf["masterplugin"] != nil {
		if !isBool(netconf["masterplugin"]) {
			return fmt.Errorf("delegate field 'masterplugin' must be a bool")
		}
	}

	if netconf["masterplugin"] != nil {
		if netconf["masterplugin"].(bool) != false && masterpluginEnabled != true {
			masterpluginEnabled = true
		} else if netconf["masterplugin"].(bool) != false && masterpluginEnabled == true {
			return fmt.Errorf("only one delegate can have 'masterplugin'")
		}
	}
	return nil
}

func isMasterplugin(netconf map[string]interface{}) bool {
	if netconf["masterplugin"] == nil {
		return false
	}

	if netconf["masterplugin"].(bool) == true {
		return true
	}

	return false
}

func delegateAdd(podif func() string, argif string, netconf map[string]interface{}, onlyMaster bool) (bool, error) {
	netconfBytes, err := json.Marshal(netconf)
	if err != nil {
		return true, fmt.Errorf("Multus: error serializing multus delegate netconf: %v", err)
	}

	if isMasterplugin(netconf) != onlyMaster {
		return true, nil
	}

	if !isMasterplugin(netconf) {
		if os.Setenv("CNI_IFNAME", podif()) != nil {
			return true, fmt.Errorf("Multus: error in setting CNI_IFNAME")
		}
	} else {
		if os.Setenv("CNI_IFNAME", argif) != nil {
			return true, fmt.Errorf("Multus: error in setting CNI_IFNAME")
		}
	}

	if netconf["ifnameRequest"] != nil {
		if os.Setenv("CNI_IFNAME", netconf["ifnameRequest"].(string)) != nil {
			return true, fmt.Errorf("Multus: error in setting CNI_IFNAME")
		}
	}

	result, err := invoke.DelegateAdd(netconf["type"].(string), netconfBytes)
	if err != nil {
		return true, fmt.Errorf("Multus: error in invoke Delegate add - %q: %v", netconf["type"].(string), err)
	}

	if !isMasterplugin(netconf) {
		return true, nil
	}

	return false, result.Print()
}

func delegateDel(podif func() string, argif string, netconf map[string]interface{}) error {
	netconfBytes, err := json.Marshal(netconf)
	if err != nil {
		return fmt.Errorf("Multus: error serializing multus delegate netconf: %v", err)
	}

	if !isMasterplugin(netconf) {
		if os.Setenv("CNI_IFNAME", podif()) != nil {
			return fmt.Errorf("Multus: error in setting CNI_IFNAME")
		}
	} else {
		if os.Setenv("CNI_IFNAME", argif) != nil {
			return fmt.Errorf("Multus: error in setting CNI_IFNAME")
		}
	}

	err = invoke.DelegateDel(netconf["type"].(string), netconfBytes)
	if err != nil {
		return fmt.Errorf("Multus: error in invoke Delegate del - %q: %v", netconf["type"].(string), err)
	}

	return err
}

func clearPlugins(mIdx int, pIdx int, argIfname string, delegates []map[string]interface{}) error {

	if os.Setenv("CNI_COMMAND", "DEL") != nil {
		return fmt.Errorf("Multus: error in setting CNI_COMMAND to DEL")
	}

	podifName := getifname()
	r := delegateDel(podifName, argIfname, delegates[mIdx])
	if r != nil {
		return r
	}

	for idx := 0; idx < pIdx && idx != mIdx; idx++ {
		r := delegateDel(podifName, argIfname, delegates[idx])
		if r != nil {
			return r
		}
	}

	return nil
}

func cmdAdd(args *skel.CmdArgs) error {
	var result error
	var nopodnet bool
	n, err := loadNetConf(args.StdinData)
	if err != nil {
		return fmt.Errorf("err in loading netconf: %v", err)
	}

	if n.Kubeconfig != "" {
		podDelegate, err := k8s.GetK8sNetwork(args, n.Kubeconfig)
		if err != nil {
			if _, ok := err.(*k8s.NoK8sNetworkError); ok {
				nopodnet = true
				if !defaultcninetwork {
					return fmt.Errorf("Multus: Err in getting k8s network from the pod spec annotation, check the pod spec or set delegate for the default network, Refer the README.md: %v", err)
				}
			} else if !defaultcninetwork {
				return fmt.Errorf("Multus: Err in getting k8s network from pod: %v", err)
			}
		}

		// If it's empty just leave it as the netconfig states (e.g. just default)
		if len(podDelegate) != 0 {
			if n.UseDefault {
				// In the case that we force the default
				// We add the found configs from CRD
				for _, eachDelegate := range podDelegate {
					eachDelegate["masterplugin"] = false
					n.Delegates = append(n.Delegates, eachDelegate)
				}

			} else {
				// Otherwise, only the CRD delegates are used.
				n.Delegates = podDelegate
			}
		}
	}

	for _, delegate := range n.Delegates {
		if err := checkDelegate(delegate); err != nil {
			return fmt.Errorf("Multus: Err in delegate conf: %v", err)
		}
	}

	if n.Kubeconfig == "" || nopodnet {
		if err := saveDelegates(args.ContainerID, n.CNIDir, n.Delegates); err != nil {
			return fmt.Errorf("Multus: Err in saving the delegates: %v", err)
		}
	}

	podifName := getifname()
	var mIndex int
	for index, delegate := range n.Delegates {
		err, r := delegateAdd(podifName, args.IfName, delegate, true)
		if err != true {
			result = r
			mIndex = index
		} else if (err != false) && r != nil {
			return r
		}
	}

	for index, delegate := range n.Delegates {
		err, r := delegateAdd(podifName, args.IfName, delegate, false)
		if err != true {
			result = r
		} else if (err != false) && r != nil {
			perr := clearPlugins(mIndex, index, args.IfName, n.Delegates)
			if perr != nil {
				return perr
			}
			return r
		}
	}

	return result
}

func cmdDel(args *skel.CmdArgs) error {
	var result error
	var nopodnet bool

	in, err := loadNetConf(args.StdinData)
	if err != nil {
		return err
	}

	if in.Kubeconfig != "" {
		podDelegate, r := k8s.GetK8sNetwork(args, in.Kubeconfig)
		if r != nil {
			if _, ok := r.(*k8s.NoK8sNetworkError); ok {
				nopodnet = true
				// no network found from default and annotaed network,
				// we do nothing to remove network for the pod!
				if !defaultcninetwork {
					return fmt.Errorf("Multus: Err in getting k8s network from the poc spec, check the pod spec or set delegate for the default network, Refer the README.md: %v", r)
				}
			} else {
				return fmt.Errorf("Multus: Err in getting k8s network from pod: %v", r)
			}
		}

		if len(podDelegate) != 0 {
			if in.UseDefault {
				// In the case that we force the default
				// We add the found configs from CRD (in reverse order)
				for i := len(podDelegate) - 1; i >= 0; i-- {
					podDelegate[i]["masterplugin"] = false
					in.Delegates = append(in.Delegates, podDelegate[i])
				}
			} else {
				in.Delegates = podDelegate
			}
		}
	}

	if in.Kubeconfig == "" || nopodnet {
		netconfBytes, err := consumeScratchNetConf(args.ContainerID, in.CNIDir)
		if err != nil {
			if os.IsNotExist(err) {
				// Per spec should ignore error if resources are missing / already removed
				return nil
			}
			return fmt.Errorf("Multus: Err in  reading the delegates: %v", err)
		}

		if err := json.Unmarshal(netconfBytes, &in.Delegates); err != nil {
			return fmt.Errorf("Multus: failed to load netconf: %v", err)
		}
	}

	podifName := getifname()
	for _, delegate := range in.Delegates {
		r := delegateDel(podifName, args.IfName, delegate)
		if r != nil {
			return r
		}
		result = r
	}

	return result
}

func main() {
	skel.PluginMain(cmdAdd, cmdDel, version.All)
}
<commit_msg>Fix default case for CNIDir<commit_after>// Copyright (c) 2017 Intel Corporation
//
// 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.

// This is a "Multi-plugin".The delegate concept refered from CNI project
// It reads other plugin netconf, and then invoke them, e.g.
// flannel or sriov plugin.

package main

import (
	"encoding/json"
	"fmt"
	"io/ioutil"
	"os"
	"path/filepath"

	k8s "github.com/Intel-Corp/multus-cni/k8sclient"
	"github.com/Intel-Corp/multus-cni/types"
	"github.com/containernetworking/cni/pkg/invoke"
	"github.com/containernetworking/cni/pkg/skel"
	"github.com/containernetworking/cni/pkg/version"
)

const defaultCNIDir = "/var/lib/cni/multus"

var masterpluginEnabled bool
var defaultcninetwork bool

//taken from cni/plugins/meta/flannel/flannel.go
func isString(i interface{}) bool {
	_, ok := i.(string)
	return ok
}

func isBool(i interface{}) bool {
	_, ok := i.(bool)
	return ok
}

func loadNetConf(bytes []byte) (*types.NetConf, error) {
	netconf := &types.NetConf{}
	if err := json.Unmarshal(bytes, netconf); err != nil {
		return nil, fmt.Errorf("failed to load netconf: %v", err)
	}

	if netconf.Kubeconfig != "" && netconf.Delegates != nil {
		defaultcninetwork = true
	}

	if netconf.CNIDir == "" {
		netconf.CNIDir = defaultCNIDir
	}

	if netconf.UseDefault {
		if netconf.Kubeconfig == "" || !defaultcninetwork {
			return nil, fmt.Errorf(`If you have set always_use_default, you must also set the delegates & the kubeconfig, refer to the README`)
		}
		return netconf, nil
	}

	if !netconf.UseDefault && (netconf.Kubeconfig != "" && !defaultcninetwork) {
		return netconf, nil
	}

	if len(netconf.Delegates) == 0 && !defaultcninetwork {
		return nil, fmt.Errorf(`delegates or kubeconfig option is must, refer README.md`)
	}

	return netconf, nil
}

func saveScratchNetConf(containerID, dataDir string, netconf []byte) error {
	if err := os.MkdirAll(dataDir, 0700); err != nil {
		return fmt.Errorf("failed to create the multus data directory(%q): %v", dataDir, err)
	}

	path := filepath.Join(dataDir, containerID)

	err := ioutil.WriteFile(path, netconf, 0600)
	if err != nil {
		return fmt.Errorf("failed to write container data in the path(%q): %v", path, err)
	}

	return err
}

func consumeScratchNetConf(containerID, dataDir string) ([]byte, error) {
	path := filepath.Join(dataDir, containerID)
	defer os.Remove(path)

	return ioutil.ReadFile(path)
}

func getifname() (f func() string) {
	var interfaceIndex int
	f = func() string {
		ifname := fmt.Sprintf("net%d", interfaceIndex)
		interfaceIndex++
		return ifname
	}

	return
}

func saveDelegates(containerID, dataDir string, delegates []map[string]interface{}) error {
	delegatesBytes, err := json.Marshal(delegates)
	if err != nil {
		return fmt.Errorf("error serializing delegate netconf: %v", err)
	}

	if err = saveScratchNetConf(containerID, dataDir, delegatesBytes); err != nil {
		return fmt.Errorf("error in saving the  delegates : %v", err)
	}

	return err
}

func checkDelegate(netconf map[string]interface{}) error {
	if netconf["type"] == nil {
		return fmt.Errorf("delegate must have the field 'type'")
	}

	if !isString(netconf["type"]) {
		return fmt.Errorf("delegate field 'type' must be a string")
	}

	if netconf["masterplugin"] != nil {
		if !isBool(netconf["masterplugin"]) {
			return fmt.Errorf("delegate field 'masterplugin' must be a bool")
		}
	}

	if netconf["masterplugin"] != nil {
		if netconf["masterplugin"].(bool) != false && masterpluginEnabled != true {
			masterpluginEnabled = true
		} else if netconf["masterplugin"].(bool) != false && masterpluginEnabled == true {
			return fmt.Errorf("only one delegate can have 'masterplugin'")
		}
	}
	return nil
}

func isMasterplugin(netconf map[string]interface{}) bool {
	if netconf["masterplugin"] == nil {
		return false
	}

	if netconf["masterplugin"].(bool) == true {
		return true
	}

	return false
}

func delegateAdd(podif func() string, argif string, netconf map[string]interface{}, onlyMaster bool) (bool, error) {
	netconfBytes, err := json.Marshal(netconf)
	if err != nil {
		return true, fmt.Errorf("Multus: error serializing multus delegate netconf: %v", err)
	}

	if isMasterplugin(netconf) != onlyMaster {
		return true, nil
	}

	if !isMasterplugin(netconf) {
		if os.Setenv("CNI_IFNAME", podif()) != nil {
			return true, fmt.Errorf("Multus: error in setting CNI_IFNAME")
		}
	} else {
		if os.Setenv("CNI_IFNAME", argif) != nil {
			return true, fmt.Errorf("Multus: error in setting CNI_IFNAME")
		}
	}

	if netconf["ifnameRequest"] != nil {
		if os.Setenv("CNI_IFNAME", netconf["ifnameRequest"].(string)) != nil {
			return true, fmt.Errorf("Multus: error in setting CNI_IFNAME")
		}
	}

	result, err := invoke.DelegateAdd(netconf["type"].(string), netconfBytes)
	if err != nil {
		return true, fmt.Errorf("Multus: error in invoke Delegate add - %q: %v", netconf["type"].(string), err)
	}

	if !isMasterplugin(netconf) {
		return true, nil
	}

	return false, result.Print()
}

func delegateDel(podif func() string, argif string, netconf map[string]interface{}) error {
	netconfBytes, err := json.Marshal(netconf)
	if err != nil {
		return fmt.Errorf("Multus: error serializing multus delegate netconf: %v", err)
	}

	if !isMasterplugin(netconf) {
		if os.Setenv("CNI_IFNAME", podif()) != nil {
			return fmt.Errorf("Multus: error in setting CNI_IFNAME")
		}
	} else {
		if os.Setenv("CNI_IFNAME", argif) != nil {
			return fmt.Errorf("Multus: error in setting CNI_IFNAME")
		}
	}

	err = invoke.DelegateDel(netconf["type"].(string), netconfBytes)
	if err != nil {
		return fmt.Errorf("Multus: error in invoke Delegate del - %q: %v", netconf["type"].(string), err)
	}

	return err
}

func clearPlugins(mIdx int, pIdx int, argIfname string, delegates []map[string]interface{}) error {

	if os.Setenv("CNI_COMMAND", "DEL") != nil {
		return fmt.Errorf("Multus: error in setting CNI_COMMAND to DEL")
	}

	podifName := getifname()
	r := delegateDel(podifName, argIfname, delegates[mIdx])
	if r != nil {
		return r
	}

	for idx := 0; idx < pIdx && idx != mIdx; idx++ {
		r := delegateDel(podifName, argIfname, delegates[idx])
		if r != nil {
			return r
		}
	}

	return nil
}

func cmdAdd(args *skel.CmdArgs) error {
	var result error
	var nopodnet bool
	n, err := loadNetConf(args.StdinData)
	if err != nil {
		return fmt.Errorf("err in loading netconf: %v", err)
	}

	if n.Kubeconfig != "" {
		podDelegate, err := k8s.GetK8sNetwork(args, n.Kubeconfig)
		if err != nil {
			if _, ok := err.(*k8s.NoK8sNetworkError); ok {
				nopodnet = true
				if !defaultcninetwork {
					return fmt.Errorf("Multus: Err in getting k8s network from the pod spec annotation, check the pod spec or set delegate for the default network, Refer the README.md: %v", err)
				}
			} else if !defaultcninetwork {
				return fmt.Errorf("Multus: Err in getting k8s network from pod: %v", err)
			}
		}

		// If it's empty just leave it as the netconfig states (e.g. just default)
		if len(podDelegate) != 0 {
			if n.UseDefault {
				// In the case that we force the default
				// We add the found configs from CRD
				for _, eachDelegate := range podDelegate {
					eachDelegate["masterplugin"] = false
					n.Delegates = append(n.Delegates, eachDelegate)
				}

			} else {
				// Otherwise, only the CRD delegates are used.
				n.Delegates = podDelegate
			}
		}
	}

	for _, delegate := range n.Delegates {
		if err := checkDelegate(delegate); err != nil {
			return fmt.Errorf("Multus: Err in delegate conf: %v", err)
		}
	}

	if n.Kubeconfig == "" || nopodnet {
		if err := saveDelegates(args.ContainerID, n.CNIDir, n.Delegates); err != nil {
			return fmt.Errorf("Multus: Err in saving the delegates: %v", err)
		}
	}

	podifName := getifname()
	var mIndex int
	for index, delegate := range n.Delegates {
		err, r := delegateAdd(podifName, args.IfName, delegate, true)
		if err != true {
			result = r
			mIndex = index
		} else if (err != false) && r != nil {
			return r
		}
	}

	for index, delegate := range n.Delegates {
		err, r := delegateAdd(podifName, args.IfName, delegate, false)
		if err != true {
			result = r
		} else if (err != false) && r != nil {
			perr := clearPlugins(mIndex, index, args.IfName, n.Delegates)
			if perr != nil {
				return perr
			}
			return r
		}
	}

	return result
}

func cmdDel(args *skel.CmdArgs) error {
	var result error
	var nopodnet bool

	in, err := loadNetConf(args.StdinData)
	if err != nil {
		return err
	}

	if in.Kubeconfig != "" {
		podDelegate, r := k8s.GetK8sNetwork(args, in.Kubeconfig)
		if r != nil {
			if _, ok := r.(*k8s.NoK8sNetworkError); ok {
				nopodnet = true
				// no network found from default and annotaed network,
				// we do nothing to remove network for the pod!
				if !defaultcninetwork {
					return fmt.Errorf("Multus: Err in getting k8s network from the poc spec, check the pod spec or set delegate for the default network, Refer the README.md: %v", r)
				}
			} else {
				return fmt.Errorf("Multus: Err in getting k8s network from pod: %v", r)
			}
		}

		if len(podDelegate) != 0 {
			if in.UseDefault {
				// In the case that we force the default
				// We add the found configs from CRD (in reverse order)
				for i := len(podDelegate) - 1; i >= 0; i-- {
					podDelegate[i]["masterplugin"] = false
					in.Delegates = append(in.Delegates, podDelegate[i])
				}
			} else {
				in.Delegates = podDelegate
			}
		}
	}

	if in.Kubeconfig == "" || nopodnet {
		netconfBytes, err := consumeScratchNetConf(args.ContainerID, in.CNIDir)
		if err != nil {
			if os.IsNotExist(err) {
				// Per spec should ignore error if resources are missing / already removed
				return nil
			}
			return fmt.Errorf("Multus: Err in  reading the delegates: %v", err)
		}

		if err := json.Unmarshal(netconfBytes, &in.Delegates); err != nil {
			return fmt.Errorf("Multus: failed to load netconf: %v", err)
		}
	}

	podifName := getifname()
	for _, delegate := range in.Delegates {
		r := delegateDel(podifName, args.IfName, delegate)
		if r != nil {
			return r
		}
		result = r
	}

	return result
}

func main() {
	skel.PluginMain(cmdAdd, cmdDel, version.All)
}
<|endoftext|>


<commit_before>package jd

import (
	"testing"
)

func TestObjectJson(t *testing.T) {
	ctx := newTestContext(t)
	checkJson(ctx, `{"a":1}`, `{"a":1}`)
	checkJson(ctx, ` { "a" : 1 } `, `{"a":1}`)
	checkJson(ctx, `{}`, `{}`)
}

func TestObjectEqual(t *testing.T) {
	ctx := newTestContext(t)
	checkEqual(ctx, `{"a":1}`, `{"a":1}`)
	checkEqual(ctx, `{"a":1}`, `{"a":1.0}`)
	checkEqual(ctx, `{"a":[1,2]}`, `{"a":[1,2]}`)
	checkEqual(ctx, `{"a":"b"}`, `{"a":"b"}`)
}

func TestObjectNotEqual(t *testing.T) {
	ctx := newTestContext(t)
	checkNotEqual(ctx, `{"a":1}`, `{"b":1}`)
	checkNotEqual(ctx, `{"a":[1,2]}`, `{"a":[2,1]}`)
	checkNotEqual(ctx, `{"a":"b"}`, `{"a":"c"}`)
}

// TODO: add unit test for object identity with setkeys metadata.
func TestObjectHash(t *testing.T) {
	ctx := newTestContext(t)
	checkHash(ctx, `{}`, `{}`, true)
	checkHash(ctx, `{"a":1}`, `{"a":1}`, true)
	checkHash(ctx, `{"a":1}`, `{"a":2}`, false)
	checkHash(ctx, `{"a":1}`, `{"b":1}`, false)
	checkHash(ctx, `{"a":1,"b":2}`, `{"b":2,"a":1}`, true)
}

func TestObjectDiff(t *testing.T) {
	ctx := newTestContext(t)
	checkDiff(ctx, `{}`, `{}`)
	checkDiff(ctx, `{"a":1}`, `{"a":1}`)
	checkDiff(ctx, `{"a":1}`, `{"a":2}`,
		`@ ["a"]`,
		`- 1`,
		`+ 2`)
	checkDiff(ctx, `{"":1}`, `{"":1}`)
	checkDiff(ctx, `{"":1}`, `{"a":2}`,
		`@ [""]`,
		`- 1`,
		`@ ["a"]`,
		`+ 2`)
	checkDiff(ctx, `{"a":{"b":{}}}`, `{"a":{"b":{"c":1},"d":2}}`,
		`@ ["a","b","c"]`,
		`+ 1`,
		`@ ["a","d"]`,
		`+ 2`)
	// regression test for issue #18
	checkDiff(ctx,
		`{"R": [{"I": [{"T": [{"V": "t","K": "N"},{"V": "T","K": "I"}]}]}]}`,
		`{"R": [{"I": [{"T": [{"V": "t","K": "N"},{"V": "Q","K": "C"},{"V": "T","K": "I"}]}]}]}`,
		`@ ["R",0,"I",0,"T",1,"K"]`,
		`- "I"`,
		`+ "C"`,
		`@ ["R",0,"I",0,"T",1,"V"]`,
		`- "T"`,
		`+ "Q"`,
		`@ ["R",0,"I",0,"T",-1]`,
		`+ {"K":"I","V":"T"}`)
}

func testObjectPatch(t *testing.T) {
	ctx := newTestContext(t)
	checkPatch(ctx, `{}`, `{}`)
	checkPatch(ctx, `{"a":1}`, `{"a":1}`)
	checkPatch(ctx, `{"a":1}`, `{"a":2}`,
		`@ ["a"]`,
		`- 1`,
		`+ 2`)
	checkPatch(ctx, `{"":1}`, `{"":1}`)
	checkPatch(ctx, `{"":1}`, `{"a":2}`,
		`@ [""]`,
		`- 1`,
		`@ ["a"]`,
		`+ 2`)
	checkPatch(ctx, `{"a":{"b":{}}}`, `{"a":{"b":{"c":1},"d":2}}`,
		`@ ["a","b","c"]`,
		`+ 1`,
		`@ ["a","d"]`,
		`+ 2`)
}

func testObjectPatchError(t *testing.T) {
	ctx := newTestContext(t)
	checkPatchError(ctx, `{}`,
		`@ ["a"]`,
		`- 1`)
	checkPatchError(ctx, `{"a":1}`,
		`@ ["a"]`,
		`+ 2`)
	checkPatchError(ctx, `{"a":1}`,
		`@ ["a"]`,
		`+ 1`)
}
<commit_msg>Convert object tests into table tests.<commit_after>package jd

import (
	"testing"
)

func TestObjectJson(t *testing.T) {
	ctx := newTestContext(t)
	tests := []struct{
		a string
		b string
	}{{
		a:`{"a":1}`,
		b: `{"a":1}`,
	}, {
		a: ` { "a" : 1 } `,
		b:`{"a":1}`,
	}, {
		a:`{}`,
		b:`{}`,
	}}

	for _, tt := range tests {
		checkJson(ctx, tt.a, tt.b)
	}
}

func TestObjectEqual(t *testing.T) {
	ctx := newTestContext(t)
	tests := []struct{
		a string
		b string
	}{{
		a:`{"a":1}`,
		b:`{"a":1}`,
	}, {
		a: `{"a":1}`,
		b:`{"a":1.0}`,
	}, {
		a: `{"a":[1,2]}`,
		b: `{"a":[1,2]}`,
	}, {
		a: `{"a":"b"}`,
		b: `{"a":"b"}`,
	}}

	for _, tt := range tests {
		checkEqual(ctx, tt.a, tt.b)
	}
}

func TestObjectNotEqual(t *testing.T) {
	ctx := newTestContext(t)
	tests := []struct{
		a string
		b string
	}{{
		a: `{"a":1}`,
		b: `{"b":1}`,
	}, {
		a: `{"a":[1,2]}`,
		b:`{"a":[2,1]}`,
	}, {
		a: `{"a":"b"}`,
		b:`{"a":"c"}`,
	}}

	for _, tt := range tests {
		checkNotEqual(ctx, tt.a, tt.b)
	}
}

// TODO: add unit test for object identity with setkeys metadata.
func TestObjectHash(t *testing.T) {
	ctx := newTestContext(t)
	tests := []struct{
		a string
		b string
		wantEqual bool
	}{{
		a: `{}`,
		b:`{}`,
		wantEqual: true,
	}, {
		a: `{"a":1}`,
		b:`{"a":1}`,
		wantEqual: true,
	}, {
		a: `{"a":1}`,
		b:`{"a":2}`,
		wantEqual: false,
	}, {
		a: `{"a":1}`,
		b:`{"b":1}`,
		wantEqual: false,
	}, {
		a: `{"a":1,"b":2}`,
		b: `{"b":2,"a":1}`,
		wantEqual: true,
	}}

	for _, tt := range tests {
		checkHash(ctx, tt.a, tt.b, tt.wantEqual)
	}
}

func TestObjectDiff(t *testing.T) {
	ctx := newTestContext(t)
	tests := []struct{
		a string
		b string
		diff []string
	}{{
		a: `{}`,
		b: `{}`,
	}, {
		a: `{"a":1}`,
		b: `{"a":1}`,
	}, {
		a: `{"a":1}`,
		b: `{"a":2}`,
		diff: ss(
			`@ ["a"]`,
			`- 1`,
			`+ 2`,
		),
	}, {
		a: `{"":1}`,
		b: `{"":1}`,
	}, {
		a: `{"":1}`,
		b: `{"a":2}`,
		diff: ss(
			`@ [""]`,
			`- 1`,
			`@ ["a"]`,
			`+ 2`,
		),
	}, {
		a: `{"a":{"b":{}}}`,
		b: `{"a":{"b":{"c":1},"d":2}}`,
		diff: ss(
			`@ ["a","b","c"]`,
			`+ 1`,
			`@ ["a","d"]`,
			`+ 2`,
		),
	}, {
		// regression test for issue #18
		a: `{"R": [{"I": [{"T": [{"V": "t","K": "N"},{"V": "T","K": "I"}]}]}]}`,
		b: `{"R": [{"I": [{"T": [{"V": "t","K": "N"},{"V": "Q","K": "C"},{"V": "T","K": "I"}]}]}]}`,
		diff: ss(
			`@ ["R",0,"I",0,"T",1,"K"]`,
			`- "I"`,
			`+ "C"`,
			`@ ["R",0,"I",0,"T",1,"V"]`,
			`- "T"`,
			`+ "Q"`,
			`@ ["R",0,"I",0,"T",-1]`,
			`+ {"K":"I","V":"T"}`,
		),
	}}

	for _, tt := range tests {
		checkDiff(ctx, tt.a, tt.b, tt.diff...)
	}
}

func testObjectPatch(t *testing.T) {
	ctx := newTestContext(t)
	tests := []struct{
		a string
		b string
		diff []string
	}{{
		a: `{}`,
		b: `{}`,
	}, {
		a: `{"a":1}`,
		b: `{"a":1}`,
	}, {
		a: `{"a":1}`,
		b: `{"a":2}`,
		diff: ss(
			`@ ["a"]`,
			`- 1`,
			`+ 2`,
		),
	}, {
		a: `{"":1}`,
		b: `{"":1}`,
	}, {
		a: `{"":1}`,
		b: `{"a":2}`,
		diff: ss(
			`@ [""]`,
			`- 1`,
			`@ ["a"]`,
			`+ 2`,
		),
	}, {
		a: `{"a":{"b":{}}}`,
		b: `{"a":{"b":{"c":1},"d":2}}`,
		diff: ss(
			`@ ["a","b","c"]`,
			`+ 1`,
			`@ ["a","d"]`,
			`+ 2`,
		),
	}}

	for _, tt := range tests {
		checkPatch(ctx, tt.a, tt.b, tt.diff...)
	}
}

func testObjectPatchError(t *testing.T) {
	ctx := newTestContext(t)
	tests := []struct{
		a string
		diff []string
	}{{
		a: `{}`,
		diff: ss(
			`@ ["a"]`,
			`- 1`,
		),
	}, {
		a: `{"a":1}`,
		diff: ss(
			`@ ["a"]`,
			`+ 2`,
		),
	}, {
		a: `{"a":1}`,
		diff: ss(
			`@ ["a"]`,
			`+ 1`,
		),
	}}

	for _, tt := range tests {
		checkPatchError(ctx, tt.a, tt.diff...)
	}
}
<|endoftext|>


<commit_before>/*
Copyright 2016 Citrix Systems, Inc

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

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

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

import (
	"fmt"
	"github.com/chiradeep/go-nitro/netscaler"
	"github.com/hashicorp/terraform/helper/resource"
	"github.com/hashicorp/terraform/terraform"
	"testing"
)

func TestAccServicegroup_basic(t *testing.T) {
	resource.Test(t, resource.TestCase{
		PreCheck:     func() { testAccPreCheck(t) },
		Providers:    testAccProviders,
		CheckDestroy: testAccCheckServicegroupDestroy,
		Steps: []resource.TestStep{
			resource.TestStep{
				Config: testAccServicegroup_basic,
				Check: resource.ComposeTestCheckFunc(
					testAccCheckServicegroupExist("netscaler_servicegroup.foo", nil),

					resource.TestCheckResourceAttr(
						"netscaler_servicegroup.foo", "servicegroupname", "test_servicegroup"),
					resource.TestCheckResourceAttr(
						"netscaler_servicegroup.foo", "servicetype", "HTTP"),
				),
			},
		},
	})
}

func testAccCheckServicegroupExist(n string, id *string) resource.TestCheckFunc {
	return func(s *terraform.State) error {
		rs, ok := s.RootModule().Resources[n]
		if !ok {
			return fmt.Errorf("Not found: %s", n)
		}

		if rs.Primary.ID == "" {
			return fmt.Errorf("No lb vserver name is set")
		}

		if id != nil {
			if *id != "" && *id != rs.Primary.ID {
				return fmt.Errorf("Resource ID has changed!")
			}

			*id = rs.Primary.ID
		}

		nsClient := testAccProvider.Meta().(*NetScalerNitroClient).client
		data, err := nsClient.FindResource(netscaler.Servicegroup.Type(), rs.Primary.ID)

		if err != nil {
			return err
		}

		if data == nil {
			return fmt.Errorf("LB vserver %s not found", n)
		}

		return nil
	}
}

func testAccCheckServicegroupDestroy(s *terraform.State) error {
	nsClient := testAccProvider.Meta().(*NetScalerNitroClient).client

	for _, rs := range s.RootModule().Resources {
		if rs.Type != "netscaler_servicegroup" {
			continue
		}

		if rs.Primary.ID == "" {
			return fmt.Errorf("No name is set")
		}

		_, err := nsClient.FindResource(netscaler.Servicegroup.Type(), rs.Primary.ID)
		if err == nil {
			return fmt.Errorf("LB vserver %s still exists", rs.Primary.ID)
		}

	}

	return nil
}

const testAccServicegroup_basic = `

resource "netscaler_lbvserver" "foo" {

  name = "foo_lb"
  ipv46 = "10.202.11.11"
  port = 80
  servicetype = "HTTP"
}

resource "netscaler_servicegroup" "foo" {

  servicegroupname = "test_servicegroup"
  servicetype = "HTTP"
  servicegroupmembers = ["172.20.0.9:80", "172.20.0.10:80", "172.20.0.11:8080"]
  lbvserver = "foo_lb"
  depends_on = ["netscaler_lbvserver.foo"]

}
`
<commit_msg>Test expects weights<commit_after>/*
Copyright 2016 Citrix Systems, Inc

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

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

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

import (
	"fmt"
	"github.com/chiradeep/go-nitro/netscaler"
	"github.com/hashicorp/terraform/helper/resource"
	"github.com/hashicorp/terraform/terraform"
	"testing"
)

func TestAccServicegroup_basic(t *testing.T) {
	resource.Test(t, resource.TestCase{
		PreCheck:     func() { testAccPreCheck(t) },
		Providers:    testAccProviders,
		CheckDestroy: testAccCheckServicegroupDestroy,
		Steps: []resource.TestStep{
			resource.TestStep{
				Config: testAccServicegroup_basic,
				Check: resource.ComposeTestCheckFunc(
					testAccCheckServicegroupExist("netscaler_servicegroup.foo", nil),

					resource.TestCheckResourceAttr(
						"netscaler_servicegroup.foo", "servicegroupname", "test_servicegroup"),
					resource.TestCheckResourceAttr(
						"netscaler_servicegroup.foo", "servicetype", "HTTP"),
				),
			},
		},
	})
}

func testAccCheckServicegroupExist(n string, id *string) resource.TestCheckFunc {
	return func(s *terraform.State) error {
		rs, ok := s.RootModule().Resources[n]
		if !ok {
			return fmt.Errorf("Not found: %s", n)
		}

		if rs.Primary.ID == "" {
			return fmt.Errorf("No lb vserver name is set")
		}

		if id != nil {
			if *id != "" && *id != rs.Primary.ID {
				return fmt.Errorf("Resource ID has changed!")
			}

			*id = rs.Primary.ID
		}

		nsClient := testAccProvider.Meta().(*NetScalerNitroClient).client
		data, err := nsClient.FindResource(netscaler.Servicegroup.Type(), rs.Primary.ID)

		if err != nil {
			return err
		}

		if data == nil {
			return fmt.Errorf("LB vserver %s not found", n)
		}

		return nil
	}
}

func testAccCheckServicegroupDestroy(s *terraform.State) error {
	nsClient := testAccProvider.Meta().(*NetScalerNitroClient).client

	for _, rs := range s.RootModule().Resources {
		if rs.Type != "netscaler_servicegroup" {
			continue
		}

		if rs.Primary.ID == "" {
			return fmt.Errorf("No name is set")
		}

		_, err := nsClient.FindResource(netscaler.Servicegroup.Type(), rs.Primary.ID)
		if err == nil {
			return fmt.Errorf("LB vserver %s still exists", rs.Primary.ID)
		}

	}

	return nil
}

const testAccServicegroup_basic = `

resource "netscaler_lbvserver" "foo" {

  name = "foo_lb"
  ipv46 = "10.202.11.11"
  port = 80
  servicetype = "HTTP"
}

resource "netscaler_servicegroup" "foo" {

  servicegroupname = "test_servicegroup"
  servicetype = "HTTP"
  servicegroupmembers = ["172.20.0.9:80:10", "172.20.0.10:80:10", "172.20.0.11:8080:20"]
  lbvserver = "foo_lb"
  depends_on = ["netscaler_lbvserver.foo"]
}
`
<|endoftext|>


<commit_before>// PromHouse
// Copyright (C) 2017 Percona LLC
//
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program. If not, see <http://www.gnu.org/licenses/>.

package storages

import (
	"context"
	"database/sql"
	"fmt"
	"strings"
	"time"

	_ "github.com/kshvakov/clickhouse" // register SQL driver
	"github.com/prometheus/client_golang/prometheus"
	"github.com/prometheus/common/model"
	"github.com/sirupsen/logrus"

	"github.com/Percona-Lab/PromHouse/util"
)

const (
	namespace = "promhouse"
	subsystem = "clickhouse"
)

// ClickHouse implements storage interface for the ClickHouse.
type ClickHouse struct {
	db       *sql.DB
	l        *logrus.Entry
	database string
	// labels   []string

	mSamplesCurrent prometheus.Gauge

	mReads      prometheus.Summary
	mReadErrors *prometheus.CounterVec

	mWrites         prometheus.Summary
	mWriteErrors    *prometheus.CounterVec
	mWrittenLabels  *prometheus.CounterVec
	mWrittenSamples prometheus.Counter
}

func NewClickHouse(dsn string, database string, init bool) (*ClickHouse, error) {
	db, err := sql.Open("clickhouse", dsn)
	if err != nil {
		return nil, err
	}
	if err = db.Ping(); err != nil {
		return nil, err
	}

	var queries []string
	if init {
		queries = append(queries, fmt.Sprintf(`DROP DATABASE IF EXISTS %s`, database))
	}
	queries = append(queries, fmt.Sprintf(`CREATE DATABASE IF NOT EXISTS %s`, database))

	// TODO use GraphiteMergeTree?
	// TODO remove fingerprint?
	// TODO add __name__, instance, job as separate columns?
	// labels := []string{"instance", "job"}
	queries = append(queries, fmt.Sprintf(`
		CREATE TABLE IF NOT EXISTS %s.samples (
			__date__ Date,
			__labels__ String,
			__fingerprint__ UInt64,
			__ts__ Int64,
			__value__ Float64
		)
		ENGINE = MergeTree(__date__, (__fingerprint__, __ts__), 8192)`, database))

	for _, q := range queries {
		if _, err = db.Exec(q); err != nil {
			return nil, err
		}
	}

	return &ClickHouse{
		db:       db,
		l:        logrus.WithField("component", "clickhouse"),
		database: database,
		// labels:   labels,

		mSamplesCurrent: prometheus.NewGauge(prometheus.GaugeOpts{
			Namespace: namespace,
			Subsystem: subsystem,
			Name:      "samples_current",
			Help:      "Current number of stored samples.",
		}),

		mReads: prometheus.NewSummary(prometheus.SummaryOpts{
			Namespace: namespace,
			Subsystem: subsystem,
			Name:      "reads",
			Help:      "Durations of successful reads.",
		}),
		mReadErrors: prometheus.NewCounterVec(prometheus.CounterOpts{
			Namespace: namespace,
			Subsystem: subsystem,
			Name:      "read_errors",
			Help:      "Number of read errors by type: canceled, other.",
		}, []string{"type"}),

		mWrites: prometheus.NewSummary(prometheus.SummaryOpts{
			Namespace: namespace,
			Subsystem: subsystem,
			Name:      "writes",
			Help:      "Durations of committed writes.",
		}),
		mWriteErrors: prometheus.NewCounterVec(prometheus.CounterOpts{
			Namespace: namespace,
			Subsystem: subsystem,
			Name:      "write_errors",
			Help:      "Number of write errors by type: canceled, other.",
		}, []string{"type"}),
		mWrittenLabels: prometheus.NewCounterVec(prometheus.CounterOpts{
			Namespace: namespace,
			Subsystem: subsystem,
			Name:      "written_labels",
			Help:      "Number of written labels by name.",
		}, []string{"name"}),
		mWrittenSamples: prometheus.NewCounter(prometheus.CounterOpts{
			Namespace: namespace,
			Subsystem: subsystem,
			Name:      "written_samples",
			Help:      "Number of written samples.",
		}),
	}, nil
}

func (ch *ClickHouse) Describe(c chan<- *prometheus.Desc) {
	ch.mSamplesCurrent.Describe(c)

	ch.mReads.Describe(c)
	ch.mReadErrors.Describe(c)

	ch.mWrites.Describe(c)
	ch.mWriteErrors.Describe(c)
	ch.mWrittenLabels.Describe(c)
	ch.mWrittenSamples.Describe(c)
}

func (ch *ClickHouse) Collect(c chan<- prometheus.Metric) {
	query := fmt.Sprintf("SELECT COUNT(*) FROM %s.samples", ch.database)
	var count uint64
	if err := ch.db.QueryRow(query).Scan(&count); err != nil {
		ch.l.Error(err)
		return
	}
	ch.mSamplesCurrent.Set(float64(count))

	ch.mSamplesCurrent.Collect(c)

	ch.mReads.Collect(c)
	ch.mReadErrors.Collect(c)

	ch.mWrites.Collect(c)
	ch.mWriteErrors.Collect(c)
	ch.mWrittenLabels.Collect(c)
	ch.mWrittenSamples.Collect(c)
}

func (ch *ClickHouse) Read(ctx context.Context, queries []Query) (data []model.Matrix, err error) {
	start := time.Now()
	defer func() {
		if err == nil {
			ch.mReads.Observe(time.Since(start).Seconds())
			return
		}

		t := "other"
		if err == context.Canceled {
			t = "canceled"
		}
		ch.mReadErrors.WithLabelValues(t).Inc()
	}()

	data = make([]model.Matrix, len(queries))
	for i, q := range queries {
		query := fmt.Sprintf("SELECT __labels__, __fingerprint__, __ts__, __value__ FROM %s.samples WHERE __ts__ >= ? AND __ts__ <= ?", ch.database)
		args := []interface{}{int64(q.Start), int64(q.End)}
		for _, matcher := range q.Matchers {
			// if matcher.Name == model.MetricNameLabel {
			// 	query += fmt.Sprintf(" AND %s", matcher.Name)
			// } else {

			switch matcher.Type {
			case MatchEqual:
				query += fmt.Sprintf(" AND visitParamExtractString(__labels__, '%s') = ?", matcher.Name)
				args = append(args, matcher.Value)
			case MatchNotEqual:
				query += fmt.Sprintf(" AND visitParamExtractString(__labels__, '%s') != ?", matcher.Name)
				args = append(args, matcher.Value)
			case MatchRegexp:
				query += fmt.Sprintf(" AND match(visitParamExtractString(__labels__, '%s'), ?)", matcher.Name)
				args = append(args, "^(?:"+matcher.Value+")$")
			case MatchNotRegexp:
				query += fmt.Sprintf(" AND NOT match(visitParamExtractString(__labels__, '%s'), ?)", matcher.Name)
				args = append(args, "^(?:"+matcher.Value+")$")
			default:
				ch.l.Panicf("unexpected match type: %d", matcher.Type)
			}
		}
		query += " ORDER BY __fingerprint__, __ts__"

		var rows *sql.Rows
		rows, err = ch.db.Query(query, args...)
		if err != nil {
			return
		}

		var ss *model.SampleStream
		var b []byte
		var fingerprint, prevFingerprint uint64
		var ts int64
		var value float64
		for rows.Next() {
			if err = rows.Scan(&b, &fingerprint, &ts, &value); err != nil {
				return
			}
			if fingerprint != prevFingerprint {
				prevFingerprint = fingerprint
				if ss != nil {
					data[i] = append(data[i], ss)
				}

				var ls model.LabelSet
				if err = ls.UnmarshalJSON(b); err != nil {
					return
				}
				ss = &model.SampleStream{
					Metric: model.Metric(ls),
				}
			}
			ss.Values = append(ss.Values, model.SamplePair{
				Timestamp: model.Time(ts),
				Value:     model.SampleValue(value),
			})
		}
		if err = rows.Close(); err != nil {
			return
		}
		if ss != nil {
			data[i] = append(data[i], ss)
		}
	}

	return
}

func (ch *ClickHouse) Write(ctx context.Context, data model.Matrix) (err error) {
	start := time.Now()
	var tx *sql.Tx
	defer func() {
		if err == nil {
			ch.mWrites.Observe(time.Since(start).Seconds())
			return
		}

		if tx != nil {
			tx.Rollback()
		}

		t := "other"
		if err == context.Canceled {
			t = "canceled"
		}
		ch.mWriteErrors.WithLabelValues(t).Inc()
	}()

	tx, err = ch.db.BeginTx(ctx, nil)
	if err != nil {
		return
	}

	// columns := 4 + len(ch.labels)
	const columns = 5
	placeholders := strings.Repeat("?, ", columns)
	query := fmt.Sprintf(
		`INSERT INTO %s.samples (__labels__, __fingerprint__, __ts__, __value__, __date__) VALUES (%s)`,
		ch.database, placeholders[:len(placeholders)-2], // cut last ", "
	)

	var stmt *sql.Stmt
	if stmt, err = tx.PrepareContext(ctx, query); err != nil {
		return err
	}
	defer func() {
		e := stmt.Close()
		if err == nil {
			err = e
		}
	}()

	for _, ss := range data {
		args := make([]interface{}, columns)

		labels := make(map[string]string, len(ss.Metric))
		for n, v := range ss.Metric {
			// if n == model.MetricNameLabel {
			// 	args[0] = string(v)
			// } else {
			labels[string(n)] = string(v)
			ch.mWrittenLabels.WithLabelValues(string(n)).Inc()
		}
		args[0] = util.MarshalLabels(labels)
		args[1] = uint64(ss.Metric.Fingerprint())

		// for i, l := range ch.labels {
		// 	args[i+1] = string(ss.Metric[model.LabelName(l)])
		// }

		for _, sp := range ss.Values {
			args[len(args)-3] = int64(sp.Timestamp)
			args[len(args)-2] = float64(sp.Value)
			args[len(args)-1] = sp.Timestamp.Time()
			_, err = stmt.ExecContext(ctx, args...)
			if err != nil {
				return
			}
		}
	}

	if err = tx.Commit(); err != nil {
		return
	}
	ch.mWrittenSamples.Add(float64(len(data)))
	return
}

// check interface
var _ Storage = (*ClickHouse)(nil)
<commit_msg>Initialise vectors.<commit_after>// PromHouse
// Copyright (C) 2017 Percona LLC
//
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program. If not, see <http://www.gnu.org/licenses/>.

package storages

import (
	"context"
	"database/sql"
	"fmt"
	"strings"
	"time"

	_ "github.com/kshvakov/clickhouse" // register SQL driver
	"github.com/prometheus/client_golang/prometheus"
	"github.com/prometheus/common/model"
	"github.com/sirupsen/logrus"

	"github.com/Percona-Lab/PromHouse/util"
)

const (
	namespace = "promhouse"
	subsystem = "clickhouse"
)

// ClickHouse implements storage interface for the ClickHouse.
type ClickHouse struct {
	db       *sql.DB
	l        *logrus.Entry
	database string
	// labels   []string

	mSamplesCurrent prometheus.Gauge

	mReads      prometheus.Summary
	mReadErrors *prometheus.CounterVec

	mWrites         prometheus.Summary
	mWriteErrors    *prometheus.CounterVec
	mWrittenLabels  *prometheus.CounterVec
	mWrittenSamples prometheus.Counter
}

func NewClickHouse(dsn string, database string, init bool) (*ClickHouse, error) {
	db, err := sql.Open("clickhouse", dsn)
	if err != nil {
		return nil, err
	}
	if err = db.Ping(); err != nil {
		return nil, err
	}

	var queries []string
	if init {
		queries = append(queries, fmt.Sprintf(`DROP DATABASE IF EXISTS %s`, database))
	}
	queries = append(queries, fmt.Sprintf(`CREATE DATABASE IF NOT EXISTS %s`, database))

	// TODO use GraphiteMergeTree?
	// TODO remove fingerprint?
	// TODO add __name__, instance, job as separate columns?
	// labels := []string{"instance", "job"}
	queries = append(queries, fmt.Sprintf(`
		CREATE TABLE IF NOT EXISTS %s.samples (
			__date__ Date,
			__labels__ String,
			__fingerprint__ UInt64,
			__ts__ Int64,
			__value__ Float64
		)
		ENGINE = MergeTree(__date__, (__fingerprint__, __ts__), 8192)`, database))

	for _, q := range queries {
		if _, err = db.Exec(q); err != nil {
			return nil, err
		}
	}

	ch := &ClickHouse{
		db:       db,
		l:        logrus.WithField("component", "clickhouse"),
		database: database,
		// labels:   labels,

		mSamplesCurrent: prometheus.NewGauge(prometheus.GaugeOpts{
			Namespace: namespace,
			Subsystem: subsystem,
			Name:      "samples_current",
			Help:      "Current number of stored samples.",
		}),

		mReads: prometheus.NewSummary(prometheus.SummaryOpts{
			Namespace: namespace,
			Subsystem: subsystem,
			Name:      "reads",
			Help:      "Durations of successful reads.",
		}),
		mReadErrors: prometheus.NewCounterVec(prometheus.CounterOpts{
			Namespace: namespace,
			Subsystem: subsystem,
			Name:      "read_errors",
			Help:      "Number of read errors by type: canceled, other.",
		}, []string{"type"}),

		mWrites: prometheus.NewSummary(prometheus.SummaryOpts{
			Namespace: namespace,
			Subsystem: subsystem,
			Name:      "writes",
			Help:      "Durations of committed writes.",
		}),
		mWriteErrors: prometheus.NewCounterVec(prometheus.CounterOpts{
			Namespace: namespace,
			Subsystem: subsystem,
			Name:      "write_errors",
			Help:      "Number of write errors by type: canceled, other.",
		}, []string{"type"}),
		mWrittenLabels: prometheus.NewCounterVec(prometheus.CounterOpts{
			Namespace: namespace,
			Subsystem: subsystem,
			Name:      "written_labels",
			Help:      "Number of written labels by name.",
		}, []string{"name"}),
		mWrittenSamples: prometheus.NewCounter(prometheus.CounterOpts{
			Namespace: namespace,
			Subsystem: subsystem,
			Name:      "written_samples",
			Help:      "Number of written samples.",
		}),
	}

	ch.mReadErrors.WithLabelValues("canceled").Set(0)
	ch.mReadErrors.WithLabelValues("other").Set(0)
	ch.mWriteErrors.WithLabelValues("canceled").Set(0)
	ch.mWriteErrors.WithLabelValues("other").Set(0)

	return ch, nil
}

func (ch *ClickHouse) Describe(c chan<- *prometheus.Desc) {
	ch.mSamplesCurrent.Describe(c)

	ch.mReads.Describe(c)
	ch.mReadErrors.Describe(c)

	ch.mWrites.Describe(c)
	ch.mWriteErrors.Describe(c)
	ch.mWrittenLabels.Describe(c)
	ch.mWrittenSamples.Describe(c)
}

func (ch *ClickHouse) Collect(c chan<- prometheus.Metric) {
	query := fmt.Sprintf("SELECT COUNT(*) FROM %s.samples", ch.database)
	var count uint64
	if err := ch.db.QueryRow(query).Scan(&count); err != nil {
		ch.l.Error(err)
		return
	}
	ch.mSamplesCurrent.Set(float64(count))

	ch.mSamplesCurrent.Collect(c)

	ch.mReads.Collect(c)
	ch.mReadErrors.Collect(c)

	ch.mWrites.Collect(c)
	ch.mWriteErrors.Collect(c)
	ch.mWrittenLabels.Collect(c)
	ch.mWrittenSamples.Collect(c)
}

func (ch *ClickHouse) Read(ctx context.Context, queries []Query) (data []model.Matrix, err error) {
	start := time.Now()
	defer func() {
		if err == nil {
			ch.mReads.Observe(time.Since(start).Seconds())
			return
		}

		t := "other"
		if err == context.Canceled {
			t = "canceled"
		}
		ch.mReadErrors.WithLabelValues(t).Inc()
	}()

	data = make([]model.Matrix, len(queries))
	for i, q := range queries {
		query := fmt.Sprintf("SELECT __labels__, __fingerprint__, __ts__, __value__ FROM %s.samples WHERE __ts__ >= ? AND __ts__ <= ?", ch.database)
		args := []interface{}{int64(q.Start), int64(q.End)}
		for _, matcher := range q.Matchers {
			// if matcher.Name == model.MetricNameLabel {
			// 	query += fmt.Sprintf(" AND %s", matcher.Name)
			// } else {

			switch matcher.Type {
			case MatchEqual:
				query += fmt.Sprintf(" AND visitParamExtractString(__labels__, '%s') = ?", matcher.Name)
				args = append(args, matcher.Value)
			case MatchNotEqual:
				query += fmt.Sprintf(" AND visitParamExtractString(__labels__, '%s') != ?", matcher.Name)
				args = append(args, matcher.Value)
			case MatchRegexp:
				query += fmt.Sprintf(" AND match(visitParamExtractString(__labels__, '%s'), ?)", matcher.Name)
				args = append(args, "^(?:"+matcher.Value+")$")
			case MatchNotRegexp:
				query += fmt.Sprintf(" AND NOT match(visitParamExtractString(__labels__, '%s'), ?)", matcher.Name)
				args = append(args, "^(?:"+matcher.Value+")$")
			default:
				ch.l.Panicf("unexpected match type: %d", matcher.Type)
			}
		}
		query += " ORDER BY __fingerprint__, __ts__"

		var rows *sql.Rows
		rows, err = ch.db.Query(query, args...)
		if err != nil {
			return
		}

		var ss *model.SampleStream
		var b []byte
		var fingerprint, prevFingerprint uint64
		var ts int64
		var value float64
		for rows.Next() {
			if err = rows.Scan(&b, &fingerprint, &ts, &value); err != nil {
				return
			}
			if fingerprint != prevFingerprint {
				prevFingerprint = fingerprint
				if ss != nil {
					data[i] = append(data[i], ss)
				}

				var ls model.LabelSet
				if err = ls.UnmarshalJSON(b); err != nil {
					return
				}
				ss = &model.SampleStream{
					Metric: model.Metric(ls),
				}
			}
			ss.Values = append(ss.Values, model.SamplePair{
				Timestamp: model.Time(ts),
				Value:     model.SampleValue(value),
			})
		}
		if err = rows.Close(); err != nil {
			return
		}
		if ss != nil {
			data[i] = append(data[i], ss)
		}
	}

	return
}

func (ch *ClickHouse) Write(ctx context.Context, data model.Matrix) (err error) {
	start := time.Now()
	var tx *sql.Tx
	defer func() {
		if err == nil {
			ch.mWrites.Observe(time.Since(start).Seconds())
			return
		}

		if tx != nil {
			tx.Rollback()
		}

		t := "other"
		if err == context.Canceled {
			t = "canceled"
		}
		ch.mWriteErrors.WithLabelValues(t).Inc()
	}()

	tx, err = ch.db.BeginTx(ctx, nil)
	if err != nil {
		return
	}

	// columns := 4 + len(ch.labels)
	const columns = 5
	placeholders := strings.Repeat("?, ", columns)
	query := fmt.Sprintf(
		`INSERT INTO %s.samples (__labels__, __fingerprint__, __ts__, __value__, __date__) VALUES (%s)`,
		ch.database, placeholders[:len(placeholders)-2], // cut last ", "
	)

	var stmt *sql.Stmt
	if stmt, err = tx.PrepareContext(ctx, query); err != nil {
		return err
	}
	defer func() {
		e := stmt.Close()
		if err == nil {
			err = e
		}
	}()

	for _, ss := range data {
		args := make([]interface{}, columns)

		labels := make(map[string]string, len(ss.Metric))
		for n, v := range ss.Metric {
			// if n == model.MetricNameLabel {
			// 	args[0] = string(v)
			// } else {
			labels[string(n)] = string(v)
			ch.mWrittenLabels.WithLabelValues(string(n)).Inc()
		}
		args[0] = util.MarshalLabels(labels)
		args[1] = uint64(ss.Metric.Fingerprint())

		// for i, l := range ch.labels {
		// 	args[i+1] = string(ss.Metric[model.LabelName(l)])
		// }

		for _, sp := range ss.Values {
			args[len(args)-3] = int64(sp.Timestamp)
			args[len(args)-2] = float64(sp.Value)
			args[len(args)-1] = sp.Timestamp.Time()
			_, err = stmt.ExecContext(ctx, args...)
			if err != nil {
				return
			}
		}
	}

	if err = tx.Commit(); err != nil {
		return
	}
	ch.mWrittenSamples.Add(float64(len(data)))
	return
}

// check interface
var _ Storage = (*ClickHouse)(nil)
<|endoftext|>


<commit_before>package netup

import (
	"archive/zip"
	"io"
	"os"
	"path"
	"path/filepath"
	"strconv"
	"strings"
)

const zipRotateCount = 5

type extractProgressor func(curr, max uint64)

func totalUncompressedSize(zr *zip.Reader) uint64 {
	var sum uint64
	for _, zf := range zr.File {
		sum += zf.UncompressedSize64
	}
	return sum
}

func newZipFileProc(dir string, stripCount int, prev, curr fileInfoTable) func(zf *zip.File) (bool, error) {
	return func(zf *zip.File) (bool, error) {
		if zf.Mode().IsDir() {
			return false, nil
		}
		zfName := stripPath(zf.Name, stripCount)
		curr[zfName] = fileInfo{
			name: zfName,
			size: zf.UncompressedSize64,
			hash: zf.CRC32,
		}
		outName := filepath.Join(dir, zfName)
		// evacuation and optimization.
		if p, ok := prev[zfName]; ok {
			r, err := p.compareWithFile(outName)
			if err != nil {
				logCompareFileFailed(err, outName)
				return false, nil
			}
			switch r {
			case fileNotMatch:
				outName = evacuateName(outName)
			case fileIsMatch:
				// skip un-changed files.
				if p.hash == zf.CRC32 {
					return false, nil
				}
			}
		}
		// rotation.
		ext := strings.ToLower(path.Ext(zfName))
		if ext == ".exe" || ext == ".dll" {
			if err := rotateFiles(outName, zipRotateCount); err != nil {
				return false, err
			}
		}
		if err := extractZipFile(zf, outName); err != nil {
			return false, err
		}
		return true, nil
	}
}

func extractZip(zipName, dir string, stripCount int, prev fileInfoTable, ep extractProgressor) (fileInfoTable, error) {
	// extract zip file.
	zr, err := zip.OpenReader(zipName)
	if err != nil {
		return nil, err
	}
	defer zr.Close()
	var (
		curr = make(fileInfoTable)
		proc = newZipFileProc(dir, stripCount, prev, curr)
		max  = totalUncompressedSize(&zr.Reader)
		sum  uint64
		sum2 uint64
	)
	defer func() {
		logInfo("extracted %d bytes", sum)
	}()
	for _, zf := range zr.File {
		extracted, err := proc(zf)
		if err != nil {
			return nil, err
		}
		if extracted {
			sum2 += zf.UncompressedSize64
		}
		sum += zf.UncompressedSize64
		if ep != nil {
			ep(sum, max)
		}
	}
	return curr, nil
}

func extractZipFile(zf *zip.File, name string) error {
	if err := os.MkdirAll(filepath.Dir(name), 0777); err != nil {
		return err
	}
	r, err := zf.Open()
	if err != nil {
		return err
	}
	defer r.Close()
	w, err := os.Create(name)
	if err != nil {
		return err
	}
	defer w.Close()
	_, err = io.Copy(w, r)
	return err
}

func stripPath(name string, count int) string {
	s := strings.Split(name, "/")
	return path.Join(s[count:]...)
}

func evacuateName(name string) string {
	base, ext := splitExt(name)
	return base + ".orig" + ext
}

func rotateFiles(name string, max int) error {
	last := rotateName(name, max)
	err := os.Remove(last)
	if err != nil && !os.IsNotExist(err) {
		return err
	}
	for i := max - 1; i >= 0; i-- {
		curr := rotateName(name, i)
		err := os.Rename(curr, last)
		if err != nil && !os.IsNotExist(err) {
			return err
		}
		last = curr
	}
	return nil
}

func rotateName(name string, index int) string {
	if index == 0 {
		return name
	}
	base, ext := splitExt(name)
	return base + "." + strconv.Itoa(index) + ext
}

func splitExt(name string) (base, ext string) {
	ext = path.Ext(name)
	base = name[:len(name)-len(ext)]
	return base, ext
}
<commit_msg>keep mtime at extract<commit_after>package netup

import (
	"archive/zip"
	"io"
	"os"
	"path"
	"path/filepath"
	"strconv"
	"strings"
)

const zipRotateCount = 5

type extractProgressor func(curr, max uint64)

func totalUncompressedSize(zr *zip.Reader) uint64 {
	var sum uint64
	for _, zf := range zr.File {
		sum += zf.UncompressedSize64
	}
	return sum
}

func newZipFileProc(dir string, stripCount int, prev, curr fileInfoTable) func(zf *zip.File) (bool, error) {
	return func(zf *zip.File) (bool, error) {
		if zf.Mode().IsDir() {
			return false, nil
		}
		zfName := stripPath(zf.Name, stripCount)
		curr[zfName] = fileInfo{
			name: zfName,
			size: zf.UncompressedSize64,
			hash: zf.CRC32,
		}
		outName := filepath.Join(dir, zfName)
		// evacuation and optimization.
		if p, ok := prev[zfName]; ok {
			r, err := p.compareWithFile(outName)
			if err != nil {
				logCompareFileFailed(err, outName)
				return false, nil
			}
			switch r {
			case fileNotMatch:
				outName = evacuateName(outName)
			case fileIsMatch:
				// skip un-changed files.
				if p.hash == zf.CRC32 {
					return false, nil
				}
			}
		}
		// rotation.
		ext := strings.ToLower(path.Ext(zfName))
		if ext == ".exe" || ext == ".dll" {
			if err := rotateFiles(outName, zipRotateCount); err != nil {
				return false, err
			}
		}
		if err := extractZipFile(zf, outName); err != nil {
			return false, err
		}
		t := zf.ModTime()
		os.Chtimes(outName, t, t)
		return true, nil
	}
}

func extractZip(zipName, dir string, stripCount int, prev fileInfoTable, ep extractProgressor) (fileInfoTable, error) {
	// extract zip file.
	zr, err := zip.OpenReader(zipName)
	if err != nil {
		return nil, err
	}
	defer zr.Close()
	var (
		curr = make(fileInfoTable)
		proc = newZipFileProc(dir, stripCount, prev, curr)
		max  = totalUncompressedSize(&zr.Reader)
		sum  uint64
		sum2 uint64
	)
	defer func() {
		logInfo("extracted %d bytes", sum)
	}()
	for _, zf := range zr.File {
		extracted, err := proc(zf)
		if err != nil {
			return nil, err
		}
		if extracted {
			sum2 += zf.UncompressedSize64
		}
		sum += zf.UncompressedSize64
		if ep != nil {
			ep(sum, max)
		}
	}
	return curr, nil
}

func extractZipFile(zf *zip.File, name string) error {
	if err := os.MkdirAll(filepath.Dir(name), 0777); err != nil {
		return err
	}
	r, err := zf.Open()
	if err != nil {
		return err
	}
	defer r.Close()
	w, err := os.Create(name)
	if err != nil {
		return err
	}
	defer w.Close()
	_, err = io.Copy(w, r)
	return err
}

func stripPath(name string, count int) string {
	s := strings.Split(name, "/")
	return path.Join(s[count:]...)
}

func evacuateName(name string) string {
	base, ext := splitExt(name)
	return base + ".orig" + ext
}

func rotateFiles(name string, max int) error {
	last := rotateName(name, max)
	err := os.Remove(last)
	if err != nil && !os.IsNotExist(err) {
		return err
	}
	for i := max - 1; i >= 0; i-- {
		curr := rotateName(name, i)
		err := os.Rename(curr, last)
		if err != nil && !os.IsNotExist(err) {
			return err
		}
		last = curr
	}
	return nil
}

func rotateName(name string, index int) string {
	if index == 0 {
		return name
	}
	base, ext := splitExt(name)
	return base + "." + strconv.Itoa(index) + ext
}

func splitExt(name string) (base, ext string) {
	ext = path.Ext(name)
	base = name[:len(name)-len(ext)]
	return base, ext
}
<|endoftext|>