<commit_before>package main

import (
  apns "github.com/anachronistic/apns"

  "gopkg.in/mgo.v2"
  "gopkg.in/mgo.v2/bson"
)

func getiOSUser() []NotificationUser {
  session, _ := mgo.Dial("127.0.0.1")
  defer session.Close()
  session.SetMode(mgo.Monotonic, true)
  db := session.DB("insapp").C("notification")
  var result []NotificationUser
  db.Find(bson.M{"os": "iOS"}).All(&result)
  return result
}

func getiOSTokenDevice() []string {
  var result []string
  notificationUsers := getiOSUser()
    for _, notif := range notificationUsers {
       result = append(result, notif.Token)
   }
   return result
}

func TriggerNotification(message string, eventId string){
  triggeriOSNotification(message, eventId)
}

func triggeriOSNotification(message string, eventId string){
  done := make(chan bool)
  devices := getiOSTokenDevice()
  for _, device := range devices {
    go sendiOSNotificationToDevice(device, message, eventId, true, done)
  }
  <- done
}

func sendiOSNotificationToDevice(token string, message string, eventId string, dev bool, done chan bool) {
  payload := apns.NewPayload()
  payload.Alert = message
  payload.Badge = 42
  payload.Sound = "bingbong.aiff"

  pn := apns.NewPushNotification()
  pn.DeviceToken = token
  pn.AddPayload(payload)
  if len(eventId) > 0 { pn.Set("id", eventId) }

  if dev {
    client := apns.NewClient("gateway.sandbox.push.apple.com:2195", "InsappDevCert.pem", "InsappDev.pem")
    client.Send(pn)
    pn.PayloadString()
  }else{
    client := apns.NewClient("gateway.push.apple.com:2195", "InsappDevProd.pem", "InsappProd.pem")
    client.Send(pn)
    pn.PayloadString()
  }

  done <- true
}
<commit_msg>Notification in production mode<commit_after>package main

import (
  apns "github.com/anachronistic/apns"

  "gopkg.in/mgo.v2"
  "gopkg.in/mgo.v2/bson"
)

func getiOSUser() []NotificationUser {
  session, _ := mgo.Dial("127.0.0.1")
  defer session.Close()
  session.SetMode(mgo.Monotonic, true)
  db := session.DB("insapp").C("notification")
  var result []NotificationUser
  db.Find(bson.M{"os": "iOS"}).All(&result)
  return result
}

func getiOSTokenDevice() []string {
  var result []string
  notificationUsers := getiOSUser()
    for _, notif := range notificationUsers {
       result = append(result, notif.Token)
   }
   return result
}

func TriggerNotification(message string, eventId string){
  triggeriOSNotification(message, eventId)
}

func triggeriOSNotification(message string, eventId string){
  done := make(chan bool)
  devices := getiOSTokenDevice()
  for _, device := range devices {
    go sendiOSNotificationToDevice(device, message, eventId, false, done)
  }
  <- done
}

func sendiOSNotificationToDevice(token string, message string, eventId string, dev bool, done chan bool) {
  payload := apns.NewPayload()
  payload.Alert = message
  payload.Badge = 42
  payload.Sound = "bingbong.aiff"

  pn := apns.NewPushNotification()
  pn.DeviceToken = token
  pn.AddPayload(payload)
  if len(eventId) > 0 { pn.Set("id", eventId) }

  if dev {
    client := apns.NewClient("gateway.sandbox.push.apple.com:2195", "InsappDevCert.pem", "InsappDev.pem")
    client.Send(pn)
    pn.PayloadString()
  }else{
    client := apns.NewClient("gateway.push.apple.com:2195", "InsappProdCert.pem", "InsappProd.pem")
    client.Send(pn)
    pn.PayloadString()
  }

  done <- true
}
<|endoftext|>


<commit_before>package scanner

import (
	"bufio"
	"bytes"
	"io"
	"math"
	"strconv"
)

// TokenChannelSize how many tokens can be buffered into the scan channel (default to 10)
var TokenChannelSize = 10

const eof = rune(0)

// Scanner scans given io.Reader for tokens
type Scanner struct {
	r       *bufio.Reader
	lastPos struct {
		line   int
		column int
	}
	pos struct {
		line   int
		column int
	}
}

// NewScanner returns a new scanner for io.Reader
func NewScanner(r io.Reader) *Scanner {
	return &Scanner{r: bufio.NewReader(r)}
}

// ScanChannel return read only channel for tokens (closes on EOF)
func (s *Scanner) ScanChannel() (token <-chan Token) {
	c := make(chan Token, TokenChannelSize)
	go func(c chan<- Token) {
		for {
			token := s.Scan()
			c <- token
			if token.Type == TokenTypeEOF {
				break
			}

		}
		close(c)
	}(c)

	return c
}

// Scan returns next token
func (s *Scanner) Scan() (token Token) {
	token.Column = s.pos.column
	token.Line = s.pos.line

	ch := s.read()

	var t TokenType
	var text string
	var val interface{}

	switch {
	case isWhitespace(ch):
		s.unread()
		t, text = s.scanWhitespace()

	case isLetter(ch):
		s.unread()
		t, text, val = s.scanIdent()

	case isNumber(ch):
		t, text, val = s.scanNumber(ch)

	case ch == '/':
		s.unread()
		t, text = s.scanComment()

	case ch == '"' || ch == '\'' || ch == '`':
		s.unread()
		t, text, val = s.scanString(ch == '`')

	case ch == '.':
		next := s.peek()
		if isNumber(next) {
			t, text, val = s.scanNumber(ch)
			break
		}
		t = TokenTypePERIOD
		text = string(ch)

	case ch == ',':
		t = TokenTypeCOMMA
		text = string(ch)

	case ch == ':':
		t = TokenTypeCOLON
		text = string(ch)

	case ch == ';':
		t = TokenTypeSEMICOLON
		text = string(ch)

	case ch == '+':
		t = TokenTypeADD
		text = string(ch)

	case ch == '-':
		t = TokenTypeSUB
		text = string(ch)

	case ch == '*':
		t = TokenTypeASTERIX
		text = string(ch)

	case ch == '&':
		t = TokenTypeAMPERSAND
		text = string(ch)

	case ch == '(':
		t = TokenTypeLPAREN
		text = string(ch)

	case ch == '[':
		t = TokenTypeLBRACK
		text = string(ch)

	case ch == '<':
		t = TokenTypeLCHEV
		text = string(ch)

	case ch == '{':
		t = TokenTypeLBRACE
		text = string(ch)

	case ch == ')':
		t = TokenTypeRPAREN
		text = string(ch)

	case ch == ']':
		t = TokenTypeRBRACK
		text = string(ch)

	case ch == '>':
		t = TokenTypeRCHEV
		text = string(ch)

	case ch == '}':
		t = TokenTypeRBRACE
		text = string(ch)

	case ch == '#':
		t = TokenTypeHASHBANG
		text = string(ch)

	case ch == '!':
		t = TokenTypeEXCL
		text = string(ch)

	case ch == '=':
		t = TokenTypeASSIGN
		text = string(ch)

	case ch == eof:
		t = TokenTypeEOF
	}

	token.Text = text
	token.Type = t
	token.Value = val

	return
}

func (s *Scanner) scanComment() (t TokenType, text string) {
	var buf bytes.Buffer

	start := s.read()
	buf.WriteRune(start)

	afterStart := s.peek()

loop:
	for {
		ch := s.read()
		switch {
		case ch == eof:
			s.unread()
			if afterStart != '/' {
				return TokenTypeUnknown, buf.String()
			}

			// Single line comment ended
			break loop
		case afterStart == '/' && ch == '\n':
			// Single line comment ended
			s.unread()
			break loop
		case afterStart == '*' && ch == '*':
			buf.WriteRune(ch)
			next := s.read()

			if next == '/' {
				// Block comment ended
				buf.WriteRune(next)
				break loop
			} else {
				s.unread()
			}
		default:
			buf.WriteRune(ch)
		}
	}

	return TokenTypeComment, buf.String()
}

func (s *Scanner) scanString(rawString bool) (TokenType, string, string) {
	var buf bytes.Buffer
	var val bytes.Buffer

	start := s.read()
	buf.WriteRune(start)

	checkRune := func(value rune) {
		// TODO handle multiple runes
		val.WriteRune(value)
	}

loop:
	for {
		ch := s.read()
		switch {
		case !rawString && ch == '\n':
			s.unread()
			// TODO s.error("Line breaks not allowed on inpreted strings")
			return TokenTypeUnknown, buf.String(), ""

		case ch == eof:
			s.unread()
			// TODO s.error
			return TokenTypeUnknown, buf.String(), ""

		case !rawString && ch == '\\':
			buf.WriteRune(ch)
			next := s.read()
			if next == eof {
				s.unread()
				continue loop
			}

			switch next {
			case start:
				val.WriteRune(next)

			case '\\':
				val.WriteRune(next)

			case '0', '1', '2', '3', '4', '5', '6', '7': // Scan octal
				s.unread()
				t, v := s.scanDigits(8, 3)
				checkRune(v)
				buf.Write(t)
				continue loop

			case 'x':
				buf.WriteRune(next)
				t, v := s.scanDigits(16, 2)
				checkRune(v)
				buf.Write(t)
				continue loop

			case 'u':
				buf.WriteRune(next)
				t, v := s.scanDigits(16, 4)
				checkRune(v)
				buf.Write(t)
				continue loop

			case 'U':
				buf.WriteRune(next)
				t, v := s.scanDigits(16, 8)
				checkRune(v)
				buf.Write(t)
				continue loop

			default:
				val.WriteRune(ch)
				val.WriteRune(next)
			}

			buf.WriteRune(next)
		default:
			buf.WriteRune(ch)
			if ch == start {
				break loop
			} else {
				val.WriteRune(ch)
			}
		}
	}

	return TokenTypeString, buf.String(), val.String()
}

func (s *Scanner) scanDigits(base, n int) ([]byte, rune) {
	var buf bytes.Buffer
	result := 0
	var ch rune
	for n > 0 {
		ch = s.read()
		digVal := digitVal(ch)

		if digVal >= base {
			s.unread()
			break
		}
		buf.WriteRune(ch)
		result += digVal * int(math.Pow(float64(base), float64(n-1)))
		n--
	}

	if n > 0 {
		// TODO	s.error("illegal char escape")
	}

	return buf.Bytes(), rune(result)
}

func (s *Scanner) scanIdent() (t TokenType, text string, val interface{}) {
	var buf bytes.Buffer
	t = TokenTypeIdent

	for {
		if ch := s.read(); ch == eof {
			s.unread()
			break
		} else if !isLetter(ch) && !isNumber(ch) {
			s.unread()
			break
		} else {
			buf.WriteRune(ch)
		}
	}

	text = buf.String()
	switch text {
	case "true", "false":
		t = TokenTypeBoolean
		val = text == "true"
	}

	return
}

func (s *Scanner) scanNumber(ch rune) (t TokenType, text string, val interface{}) {
	var buf bytes.Buffer
	t = TokenTypeNumber

loop:
	for {
		switch {
		case isNumber(ch):
			buf.WriteRune(ch)
		case ch == '.' && t == TokenTypeNumber:
			t = TokenTypeFloat
			buf.WriteRune(ch)
		default:
			s.unread()
			break loop
		}

		ch = s.read()
	}

	text = buf.String()

	if t == TokenTypeNumber {
		// TODO process error
		val, _ = strconv.ParseInt(text, 10, 64)
	} else {
		// TODO process error
		val, _ = strconv.ParseFloat(text, 64)
	}

	return t, text, val
}

func (s *Scanner) scanWhitespace() (t TokenType, text string) {
	var buf bytes.Buffer

	for {
		if ch := s.read(); ch == eof || !isWhitespace(ch) {
			s.unread()
			break
		} else {
			buf.WriteRune(ch)
		}
	}

	return TokenTypeWhitespace, buf.String()
}

func (s *Scanner) read() rune {
	s.lastPos = s.pos
	ch, _, err := s.r.ReadRune()
	if err != nil {
		return eof
	}
	if ch == '\n' {
		s.pos.column = 0
		s.pos.line++
	} else {
		s.pos.column++
	}
	return ch
}

func (s *Scanner) peek() rune {
	defer s.unread()
	return s.read()
}

// unread places the previously read rune back on the reader.
func (s *Scanner) unread() {
	s.pos = s.lastPos
	_ = s.r.UnreadRune()
}

func isLetter(ch rune) bool {
	return (ch >= 'a' && ch <= 'z') || (ch >= 'A' && ch <= 'Z')
}

func isNumber(ch rune) bool {
	return ch >= '0' && ch <= '9'
}

func isWhitespace(ch rune) bool {
	return ch == ' ' || ch == '\t' || ch == '\n'
}

func digitVal(ch rune) int {
	switch {
	case '0' <= ch && ch <= '9':
		return int(ch - '0')
	case 'a' <= ch && ch <= 'f':
		return int(ch - 'a' + 10)
	case 'A' <= ch && ch <= 'F':
		return int(ch - 'A' + 10)
	}
	return 16
}
<commit_msg>call s.error on errors<commit_after>package scanner

import (
	"bufio"
	"bytes"
	"io"
	"math"
	"strconv"
)

// TokenChannelSize how many tokens can be buffered into the scan channel (default to 10)
var TokenChannelSize = 10

const eof = rune(0)

// Scanner scans given io.Reader for tokens
type Scanner struct {
	r       *bufio.Reader
	lastPos struct {
		line   int
		column int
	}
	pos struct {
		line   int
		column int
	}
}

// NewScanner returns a new scanner for io.Reader
func NewScanner(r io.Reader) *Scanner {
	return &Scanner{r: bufio.NewReader(r)}
}

// ScanChannel return read only channel for tokens (closes on EOF)
func (s *Scanner) ScanChannel() (token <-chan Token) {
	c := make(chan Token, TokenChannelSize)
	go func(c chan<- Token) {
		for {
			token := s.Scan()
			c <- token
			if token.Type == TokenTypeEOF {
				break
			}

		}
		close(c)
	}(c)

	return c
}

// Scan returns next token
func (s *Scanner) Scan() (token Token) {
	token.Column = s.pos.column
	token.Line = s.pos.line

	ch := s.read()

	var t TokenType
	var text string
	var val interface{}

	switch {
	case isWhitespace(ch):
		s.unread()
		t, text = s.scanWhitespace()

	case isLetter(ch):
		s.unread()
		t, text, val = s.scanIdent()

	case isNumber(ch):
		t, text, val = s.scanNumber(ch)

	case ch == '/':
		s.unread()
		t, text = s.scanComment()

	case ch == '"' || ch == '\'' || ch == '`':
		s.unread()
		t, text, val = s.scanString(ch == '`')

	case ch == '.':
		next := s.peek()
		if isNumber(next) {
			t, text, val = s.scanNumber(ch)
			break
		}
		t = TokenTypePERIOD
		text = string(ch)

	case ch == ',':
		t = TokenTypeCOMMA
		text = string(ch)

	case ch == ':':
		t = TokenTypeCOLON
		text = string(ch)

	case ch == ';':
		t = TokenTypeSEMICOLON
		text = string(ch)

	case ch == '+':
		t = TokenTypeADD
		text = string(ch)

	case ch == '-':
		t = TokenTypeSUB
		text = string(ch)

	case ch == '*':
		t = TokenTypeASTERIX
		text = string(ch)

	case ch == '&':
		t = TokenTypeAMPERSAND
		text = string(ch)

	case ch == '(':
		t = TokenTypeLPAREN
		text = string(ch)

	case ch == '[':
		t = TokenTypeLBRACK
		text = string(ch)

	case ch == '<':
		t = TokenTypeLCHEV
		text = string(ch)

	case ch == '{':
		t = TokenTypeLBRACE
		text = string(ch)

	case ch == ')':
		t = TokenTypeRPAREN
		text = string(ch)

	case ch == ']':
		t = TokenTypeRBRACK
		text = string(ch)

	case ch == '>':
		t = TokenTypeRCHEV
		text = string(ch)

	case ch == '}':
		t = TokenTypeRBRACE
		text = string(ch)

	case ch == '#':
		t = TokenTypeHASHBANG
		text = string(ch)

	case ch == '!':
		t = TokenTypeEXCL
		text = string(ch)

	case ch == '=':
		t = TokenTypeASSIGN
		text = string(ch)

	case ch == eof:
		t = TokenTypeEOF
	}

	token.Text = text
	token.Type = t
	token.Value = val

	return
}

func (s *Scanner) scanComment() (t TokenType, text string) {
	var buf bytes.Buffer

	start := s.read()
	buf.WriteRune(start)

	afterStart := s.peek()

loop:
	for {
		ch := s.read()
		switch {
		case ch == eof:
			s.unread()
			if afterStart != '/' {
				return TokenTypeUnknown, buf.String()
			}

			// Single line comment ended
			break loop
		case afterStart == '/' && ch == '\n':
			// Single line comment ended
			s.unread()
			break loop
		case afterStart == '*' && ch == '*':
			buf.WriteRune(ch)
			next := s.read()

			if next == '/' {
				// Block comment ended
				buf.WriteRune(next)
				break loop
			} else {
				s.unread()
			}
		default:
			buf.WriteRune(ch)
		}
	}

	return TokenTypeComment, buf.String()
}

func (s *Scanner) scanString(rawString bool) (TokenType, string, string) {
	var buf bytes.Buffer
	var val bytes.Buffer

	start := s.read()
	buf.WriteRune(start)

	checkRune := func(value rune) {
		// TODO handle multiple runes
		val.WriteRune(value)
	}

loop:
	for {
		ch := s.read()
		switch {
		case !rawString && ch == '\n':
			s.unread()
			s.error("Line breaks not allowed on intepreted string literals")
			return TokenTypeUnknown, buf.String(), ""

		case ch == eof:
			s.unread()
			s.error("EOF before string closed")
			return TokenTypeUnknown, buf.String(), ""

		case !rawString && ch == '\\':
			buf.WriteRune(ch)
			next := s.read()
			if next == eof {
				s.unread()
				continue loop
			}

			switch next {
			case start:
				val.WriteRune(next)

			case '\\':
				val.WriteRune(next)

			case '0', '1', '2', '3', '4', '5', '6', '7': // Scan octal
				s.unread()
				t, v := s.scanDigits(8, 3)
				checkRune(v)
				buf.Write(t)
				continue loop

			case 'x':
				buf.WriteRune(next)
				t, v := s.scanDigits(16, 2)
				checkRune(v)
				buf.Write(t)
				continue loop

			case 'u':
				buf.WriteRune(next)
				t, v := s.scanDigits(16, 4)
				checkRune(v)
				buf.Write(t)
				continue loop

			case 'U':
				buf.WriteRune(next)
				t, v := s.scanDigits(16, 8)
				checkRune(v)
				buf.Write(t)
				continue loop

			default:
				val.WriteRune(ch)
				val.WriteRune(next)
			}

			buf.WriteRune(next)
		default:
			buf.WriteRune(ch)
			if ch == start {
				break loop
			} else {
				val.WriteRune(ch)
			}
		}
	}

	return TokenTypeString, buf.String(), val.String()
}

func (s *Scanner) scanDigits(base, n int) ([]byte, rune) {
	var buf bytes.Buffer
	result := 0
	var ch rune
	for n > 0 {
		ch = s.read()
		digVal := digitVal(ch)

		if digVal >= base {
			s.unread()
			break
		}
		buf.WriteRune(ch)
		result += digVal * int(math.Pow(float64(base), float64(n-1)))
		n--
	}

	if n > 0 {
		s.error("illegal char escape")
	}

	return buf.Bytes(), rune(result)
}

func (s *Scanner) scanIdent() (t TokenType, text string, val interface{}) {
	var buf bytes.Buffer
	t = TokenTypeIdent

	for {
		if ch := s.read(); ch == eof {
			s.unread()
			break
		} else if !isLetter(ch) && !isNumber(ch) {
			s.unread()
			break
		} else {
			buf.WriteRune(ch)
		}
	}

	text = buf.String()
	switch text {
	case "true", "false":
		t = TokenTypeBoolean
		val = text == "true"
	}

	return
}

func (s *Scanner) scanNumber(ch rune) (t TokenType, text string, val interface{}) {
	var buf bytes.Buffer
	t = TokenTypeNumber

loop:
	for {
		switch {
		case isNumber(ch):
			buf.WriteRune(ch)
		case ch == '.' && t == TokenTypeNumber:
			t = TokenTypeFloat
			buf.WriteRune(ch)
		default:
			s.unread()
			break loop
		}

		ch = s.read()
	}

	text = buf.String()

	var err error
	if t == TokenTypeNumber {
		val, err = strconv.ParseInt(text, 10, 64)
		if err != nil {
			s.error(err.Error())
		}
	} else {
		val, err = strconv.ParseFloat(text, 64)
		if err != nil {
			s.error(err.Error())
		}
	}

	return t, text, val
}

func (s *Scanner) scanWhitespace() (t TokenType, text string) {
	var buf bytes.Buffer

	for {
		if ch := s.read(); ch == eof || !isWhitespace(ch) {
			s.unread()
			break
		} else {
			buf.WriteRune(ch)
		}
	}

	return TokenTypeWhitespace, buf.String()
}

func (s *Scanner) read() rune {
	s.lastPos = s.pos
	ch, _, err := s.r.ReadRune()
	if err != nil {
		return eof
	}
	if ch == '\n' {
		s.pos.column = 0
		s.pos.line++
	} else {
		s.pos.column++
	}
	return ch
}

func (s *Scanner) error(err string) {
	// TODO process error
}

func (s *Scanner) peek() rune {
	defer s.unread()
	return s.read()
}

// unread places the previously read rune back on the reader.
func (s *Scanner) unread() {
	s.pos = s.lastPos
	_ = s.r.UnreadRune()
}

func isLetter(ch rune) bool {
	return (ch >= 'a' && ch <= 'z') || (ch >= 'A' && ch <= 'Z')
}

func isNumber(ch rune) bool {
	return ch >= '0' && ch <= '9'
}

func isWhitespace(ch rune) bool {
	return ch == ' ' || ch == '\t' || ch == '\n'
}

func digitVal(ch rune) int {
	switch {
	case '0' <= ch && ch <= '9':
		return int(ch - '0')
	case 'a' <= ch && ch <= 'f':
		return int(ch - 'a' + 10)
	case 'A' <= ch && ch <= 'F':
		return int(ch - 'A' + 10)
	}
	return 16
}
<|endoftext|>


<commit_before>package models

import (
	// Go Builtin Packages
	"bytes"
	"compress/zlib"
	"crypto/sha256"
	"encoding/hex"
	"log"
	"net"
	"net/http"
	"strings"
	"time"

	// Google Appengine Packages
	"appengine"
	"appengine/datastore"
	"appengine/user"
)

type Tags []string
type Paste struct {
	UserID  string    `datastore:"user_id"`
	Title   string    `datastore:"title"`
	Content []byte    `datastore:"content,noindex"`
	Tags    Tags      `datastore:"tags"`
	Format  string    `datastore:"format,noindex"`
	IPAddr  net.IP    `datastore:"ipaddr,noindex"`
	Date    time.Time `datastore:"date_published"`
	// We need the Zlib flag to correctly process old, uncompressed content
	Zlib bool `datastore:"zlib,noindex"`
}

func (p *Paste) Load(ds <-chan datastore.Property) error {
	// TODO: Do something with ErrFieldMismatch here
	if err := datastore.LoadStruct(p, ds); err != nil {
		return nil // Do nothing D:
	}
	return nil
}

func (p *Paste) Save(ds chan<- datastore.Property) error {
	return datastore.SaveStruct(p, ds)
}

const PasteDSKind string = "Paste"

func genpasteKey(c appengine.Context, p *Paste) (*datastore.Key, string) {
	timestamp := time.Now().Format(time.StampNano)

	hasher := sha256.New()
	hasher.Write([]byte(timestamp))
	hasher.Write(p.Content)
	digest := hex.EncodeToString(hasher.Sum(nil))

	paste_id := digest[:8] // This is probably a silly way to go about it xD
	return datastore.NewKey(c, PasteDSKind, paste_id, 0, nil), paste_id
}

func (p Paste) validate() error {
	// FIXME: Implement input validation here
	return nil
}

func (p Paste) save(c appengine.Context) (string, error) {
	if err := p.validate(); err == nil {
		key, stringID := genpasteKey(c, &p)
		_, err := datastore.Put(c, key, &p)
		if err != nil {
			log.Panic(err)
		}
		return stringID, nil
	} else {
		return "", err
	}
}

func (p Paste) Delete(c appengine.Context, paste_id string) {
	key := datastore.NewKey(c, PasteDSKind, paste_id, 0, nil)
	if err := datastore.Delete(c, key); err != nil {
		log.Panic(err)
	}
}

func NewPaste(c appengine.Context, r *http.Request) string {
	var paste Paste

	if usr := user.Current(c); usr != nil {
		paste.UserID = usr.ID
	}

	paste.Title = r.PostForm.Get("title")

	var content bytes.Buffer
	w := zlib.NewWriter(&content)
	w.Write([]byte(r.PostForm.Get("content")))
	w.Close()
	paste.Content = content.Bytes()
	paste.Zlib = true

	paste.Tags = strings.Split(r.PostForm.Get("tags"), " ")
	paste.Format = r.PostForm.Get("format")

	if ipaddr := net.ParseIP(r.RemoteAddr); ipaddr != nil {
		paste.IPAddr = net.IP(ipaddr)
	}

	paste.Date = time.Now()

	stringID, _ := paste.save(c) // FIXME: do something if this returns an error
	// stringID := "meep" // DEBUG: Let's not write to the datastore at the moment :o
	return stringID
}

func GetPaste(c appengine.Context, paste_id string) (*Paste, error) {
	key := datastore.NewKey(c, PasteDSKind, paste_id, 0, nil)
	paste := new(Paste)
	if err := datastore.Get(c, key, paste); err != nil {
		return paste, err
	}
	return paste, nil
}
<commit_msg>Ohey, we already had the current time<commit_after>package models

import (
	// Go Builtin Packages
	"bytes"
	"compress/zlib"
	"crypto/sha256"
	"encoding/hex"
	"log"
	"net"
	"net/http"
	"strings"
	"time"

	// Google Appengine Packages
	"appengine"
	"appengine/datastore"
	"appengine/user"
)

type Tags []string
type Paste struct {
	UserID  string    `datastore:"user_id"`
	Title   string    `datastore:"title"`
	Content []byte    `datastore:"content,noindex"`
	Tags    Tags      `datastore:"tags"`
	Format  string    `datastore:"format,noindex"`
	IPAddr  net.IP    `datastore:"ipaddr,noindex"`
	Date    time.Time `datastore:"date_published"`
	// We need the Zlib flag to correctly process old, uncompressed content
	Zlib bool `datastore:"zlib,noindex"`
}

func (p *Paste) Load(ds <-chan datastore.Property) error {
	// TODO: Do something with ErrFieldMismatch here
	if err := datastore.LoadStruct(p, ds); err != nil {
		return nil // Do nothing D:
	}
	return nil
}

func (p *Paste) Save(ds chan<- datastore.Property) error {
	return datastore.SaveStruct(p, ds)
}

const PasteDSKind string = "Paste"

func genpasteKey(c appengine.Context, p *Paste) (*datastore.Key, string) {
	timestamp := p.Date.Format(time.StampNano)

	hasher := sha256.New()
	hasher.Write([]byte(timestamp))
	hasher.Write(p.Content)
	digest := hex.EncodeToString(hasher.Sum(nil))

	paste_id := digest[:8] // This is probably a silly way to go about it xD
	return datastore.NewKey(c, PasteDSKind, paste_id, 0, nil), paste_id
}

func (p Paste) validate() error {
	// FIXME: Implement input validation here
	return nil
}

func (p Paste) save(c appengine.Context) (string, error) {
	if err := p.validate(); err == nil {
		key, stringID := genpasteKey(c, &p)
		_, err := datastore.Put(c, key, &p)
		if err != nil {
			log.Panic(err)
		}
		return stringID, nil
	} else {
		return "", err
	}
}

func (p Paste) Delete(c appengine.Context, paste_id string) {
	key := datastore.NewKey(c, PasteDSKind, paste_id, 0, nil)
	if err := datastore.Delete(c, key); err != nil {
		log.Panic(err)
	}
}

func NewPaste(c appengine.Context, r *http.Request) string {
	var paste Paste

	if usr := user.Current(c); usr != nil {
		paste.UserID = usr.ID
	}

	paste.Title = r.PostForm.Get("title")

	var content bytes.Buffer
	w := zlib.NewWriter(&content)
	w.Write([]byte(r.PostForm.Get("content")))
	w.Close()
	paste.Content = content.Bytes()
	paste.Zlib = true

	paste.Tags = strings.Split(r.PostForm.Get("tags"), " ")
	paste.Format = r.PostForm.Get("format")

	if ipaddr := net.ParseIP(r.RemoteAddr); ipaddr != nil {
		paste.IPAddr = net.IP(ipaddr)
	}

	paste.Date = time.Now()

	stringID, _ := paste.save(c) // FIXME: do something if this returns an error
	// stringID := "meep" // DEBUG: Let's not write to the datastore at the moment :o
	return stringID
}

func GetPaste(c appengine.Context, paste_id string) (*Paste, error) {
	key := datastore.NewKey(c, PasteDSKind, paste_id, 0, nil)
	paste := new(Paste)
	if err := datastore.Get(c, key, paste); err != nil {
		return paste, err
	}
	return paste, nil
}
<|endoftext|>


<commit_before>////////////////////////////////////////////////////////////////////////////
// Purpose: Turn the duoyinzi-phrase.txt file to json as Go dictionary
// Authors: Tong Sun (c) 2017
// Sources: https://github.com/mozillazg/phrase-pinyin-data
////////////////////////////////////////////////////////////////////////////

package main

import (
	"bytes"
	"fmt"
	"os"
	"strings"

	"github.com/spakin/awk"
)

// json string to return
var bufRet = bytes.NewBufferString("")

func main() {
	// https://godoc.org/github.com/spakin/awk
	s := awk.NewScript()

	// == BEGIN
	sa := s.NewValueArray()
	// last word, and its length
	sa.Set("wLast", "")
	sa.Set("cLength", 0)
	bufRet.WriteByte('{')

	// == Match & Process
	s.AppendStmt(nil, func(s *awk.Script) {
		ww, py := s.F(1).String(), ""
		for ii := 2; ii <= s.NF; ii++ {
			py += s.F(ii).String() + " "
		}
		ww = ww[:len(ww)-1] // last char is ":"
		print(s.NR, ww, py)

		// count of current word length and match length with last word
		cLength := sa.Get("cLength").Int()
		cMatch := commPrefixLen(sa.Get("wLast").String(), ww)
		print(" ", cLength, " ", cMatch)

		lDiff := cLength - cMatch
		if lDiff == 0 && sa.Get("wLast").String() != "" {
			// the new phrase is longer than last one, ignore it
			print("\n")
			s.Next()
		}

		sa.Set("wLast", ww)
		//print(" Saved:", sa.Get("wLast").String())
		// Only partial match, close the json, with (lDiff-1) closes
		for ii := 1; ii < lDiff; ii++ {
			fmt.Fprintf(bufRet, "},")
		}
		outputEntry(ww, py, cMatch)
		sa.Set("cLength", len([]rune(ww)))
	})

	// == END
	s.End = func(s *awk.Script) {
		for ii := 0; ii < sa.Get("cLength").Int(); ii++ {
			bufRet.WriteByte('}')
		}
		//ret := bufRet.String()
		ret := strings.Replace(bufRet.String(), ",}", "}", -1)
		fmt.Printf("%s\n", ret)
	}

	if err := s.Run(os.Stdin); err != nil {
		panic(err)
	}
}

func outputEntry(s, py string, start int) {
	print(" output ", s, " from ", start, "\n")
	rs := []rune(s)

	for ii := start; ii < len(rs)-1; ii++ {
		fmt.Fprintf(bufRet, `"%s":{`, string(rs[ii]))
	}
	fmt.Fprintf(bufRet, `"%s":"%s",`, string(rs[len(rs)-1]), py)
}

func commPrefixLen(s1, s2 string) int {
	print(" compare ", s1, ":", s2)
	rs1, rs2 := []rune(s1), []rune(s2)
	ii := 0
	for ; ii < min(len(rs1), len(rs2)); ii++ {
		if rs1[ii] != rs2[ii] {
			return ii
		}
	}
	return ii
}

func min(x, y int) int {
	if x < y {
		return x
	}
	return y
}

func max(x, y int) int {
	if x > y {
		return x
	}
	return y
}

func debug(args ...interface{}) {

}

/*
 */
<commit_msg>- [*] disable duoyinzi-dict-gen.go debugging<commit_after>////////////////////////////////////////////////////////////////////////////
// Purpose: Turn the duoyinzi-phrase.txt file to json as Go dictionary
// Authors: Tong Sun (c) 2017
// Sources: https://github.com/mozillazg/phrase-pinyin-data
////////////////////////////////////////////////////////////////////////////

package main

import (
	"bytes"
	"fmt"
	"os"
	"strings"

	"github.com/spakin/awk"
)

// json string to return
var bufRet = bytes.NewBufferString("")

// var p = print X: use of builtin print not in function call
// Remove "//D>" to debug

func main() {
	// https://godoc.org/github.com/spakin/awk
	s := awk.NewScript()

	// == BEGIN
	sa := s.NewValueArray()
	// last word, and its length
	sa.Set("wLast", "")
	sa.Set("cLength", 0)
	bufRet.WriteByte('{')

	// == Match & Process
	s.AppendStmt(nil, func(s *awk.Script) {
		ww, py := s.F(1).String(), ""
		for ii := 2; ii <= s.NF; ii++ {
			py += s.F(ii).String() + " "
		}
		ww = ww[:len(ww)-1] // last char is ":"
		//D>print(s.NR, ww, py)

		// count of current word length and match length with last word
		cLength := sa.Get("cLength").Int()
		cMatch := commPrefixLen(sa.Get("wLast").String(), ww)
		//D>print(" ", cLength, " ", cMatch)

		lDiff := cLength - cMatch
		if lDiff == 0 && sa.Get("wLast").String() != "" {
			// the new phrase is longer than last one, ignore it
			//D>println()
			s.Next()
		}

		sa.Set("wLast", ww)
		//print(" Saved:", sa.Get("wLast").String())
		// Only partial match, close the json, with (lDiff-1) closes
		for ii := 1; ii < lDiff; ii++ {
			fmt.Fprintf(bufRet, "},")
		}
		outputEntry(ww, py, cMatch)
		sa.Set("cLength", len([]rune(ww)))
	})

	// == END
	s.End = func(s *awk.Script) {
		for ii := 0; ii < sa.Get("cLength").Int(); ii++ {
			bufRet.WriteByte('}')
		}
		//ret := bufRet.String()
		ret := strings.Replace(bufRet.String(), ",}", "}", -1)
		fmt.Printf("%s\n", ret)
	}

	if err := s.Run(os.Stdin); err != nil {
		panic(err)
	}
}

func outputEntry(s, py string, start int) {
	//D>println(" output ", s, " from ", start)
	rs := []rune(s)

	for ii := start; ii < len(rs)-1; ii++ {
		fmt.Fprintf(bufRet, `"%s":{`, string(rs[ii]))
	}
	fmt.Fprintf(bufRet, `"%s":"%s",`, string(rs[len(rs)-1]), py)
}

func commPrefixLen(s1, s2 string) int {
	//D>print(" compare ", s1, ":", s2)
	rs1, rs2 := []rune(s1), []rune(s2)
	ii := 0
	for ; ii < min(len(rs1), len(rs2)); ii++ {
		if rs1[ii] != rs2[ii] {
			return ii
		}
	}
	return ii
}

func min(x, y int) int {
	if x < y {
		return x
	}
	return y
}

func max(x, y int) int {
	if x > y {
		return x
	}
	return y
}

func debug(args ...interface{}) {
	// print(args...) X: invalid use of ... with builtin print
}

/*
 */
<|endoftext|>


<commit_before>package logger

import (
	log "github.com/sirupsen/logrus"
)

type NullWriter struct {
}

func (w *NullWriter) Write(p []byte) (n int, err error) {
	return 0, nil
}

type JsonLogger struct {
	logger *log.Logger
}

func NewJsonLogger(hidden bool) *JsonLogger {
	logger := &Logger{}
	logger.logger = log.New()

	if hidden == true {
		logger.logger.Out = &NullWriter{}
	}

	logger.logger.Formatter = &log.JSONFormatter{
		FieldMap: log.FieldMap{
			log.FieldKeyTime:  "@timestamp",
			log.FieldKeyLevel: "@level",
			log.FieldKeyMsg:   "@message",
		},
	}

	return logger
}

func (logger *JsonLogger) WithFields(fields map[string]interface{}) *log.Entry {
	f := make(log.Fields)
	for k, v := range fields {
		f[k] = v
	}
	return logger.logger.WithFields(fields)
}
<commit_msg>added logger for json output<commit_after>package logger

import (
	log "github.com/sirupsen/logrus"
)

// NullWriter does not write log
type NullWriter struct {
}

// Write write nothging
func (w *NullWriter) Write(p []byte) (n int, err error) {
	return 0, nil
}

// JsonLogger is a logger JSON formatter wrap
type JsonLogger struct {
	logger *log.Logger
}

// NewJsonLogger returns a JsonLogger
func NewJsonLogger(hidden bool) *JsonLogger {
	logger := &Logger{}
	logger.logger = log.New()

	if hidden == true {
		logger.logger.Out = &NullWriter{}
	}

	logger.logger.Formatter = &log.JSONFormatter{
		FieldMap: log.FieldMap{
			log.FieldKeyTime:  "@timestamp",
			log.FieldKeyLevel: "@level",
			log.FieldKeyMsg:   "@message",
		},
	}

	return logger
}

// WithFields returns a log entry
func (logger *JsonLogger) WithFields(fields map[string]interface{}) *log.Entry {
	f := make(log.Fields)
	for k, v := range fields {
		f[k] = v
	}
	return logger.logger.WithFields(fields)
}
<|endoftext|>


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

package agent

import (
	"errors"
	"fmt"
	"net/url"
	"os"
	"os/exec"
	"syscall"
	"time"

	"github.com/coreos/etcd/pkg/fileutil"
	"github.com/coreos/etcd/pkg/proxy"
	"github.com/coreos/etcd/tools/functional-tester/rpcpb"

	"go.uber.org/zap"
)

// return error for system errors (e.g. fail to create files)
// return status error in response for wrong configuration/operation (e.g. start etcd twice)
func (srv *Server) handleTesterRequest(req *rpcpb.Request) (resp *rpcpb.Response, err error) {
	defer func() {
		if err == nil {
			srv.last = req.Operation
			srv.lg.Info("handler success", zap.String("operation", req.Operation.String()))
		}
	}()

	switch req.Operation {
	case rpcpb.Operation_InitialStartEtcd:
		return srv.handleInitialStartEtcd(req)
	case rpcpb.Operation_RestartEtcd:
		return srv.handleRestartEtcd()
	case rpcpb.Operation_KillEtcd:
		return srv.handleKillEtcd()
	case rpcpb.Operation_FailArchive:
		return srv.handleFailArchive()
	case rpcpb.Operation_DestroyEtcdAgent:
		return srv.handleDestroyEtcdAgent()

	case rpcpb.Operation_BlackholePeerPortTxRx:
		return srv.handleBlackholePeerPortTxRx()
	case rpcpb.Operation_UnblackholePeerPortTxRx:
		return srv.handleUnblackholePeerPortTxRx()
	case rpcpb.Operation_DelayPeerPortTxRx:
		return srv.handleDelayPeerPortTxRx()
	case rpcpb.Operation_UndelayPeerPortTxRx:
		return srv.handleUndelayPeerPortTxRx()

	default:
		msg := fmt.Sprintf("operation not found (%v)", req.Operation)
		return &rpcpb.Response{Success: false, Status: msg}, errors.New(msg)
	}
}

func (srv *Server) handleInitialStartEtcd(req *rpcpb.Request) (*rpcpb.Response, error) {
	if srv.last != rpcpb.Operation_NotStarted {
		return &rpcpb.Response{
			Success: false,
			Status:  fmt.Sprintf("%q is not valid; last server operation was %q", rpcpb.Operation_InitialStartEtcd.String(), srv.last.String()),
		}, nil
	}

	srv.Member = req.Member
	srv.Tester = req.Tester

	err := fileutil.TouchDirAll(srv.Member.BaseDir)
	if err != nil {
		return nil, err
	}
	srv.lg.Info("created base directory", zap.String("path", srv.Member.BaseDir))

	if err = srv.createEtcdFile(); err != nil {
		return nil, err
	}
	srv.creatEtcdCmd()

	err = srv.startEtcdCmd()
	if err != nil {
		return nil, err
	}
	srv.lg.Info("started etcd", zap.String("command-path", srv.etcdCmd.Path))

	// wait some time for etcd listener start
	// before setting up proxy
	time.Sleep(time.Second)
	if err = srv.startProxy(); err != nil {
		return nil, err
	}

	return &rpcpb.Response{
		Success: true,
		Status:  "successfully started etcd!",
	}, nil
}

func (srv *Server) startProxy() error {
	if srv.Member.EtcdClientProxy {
		advertiseClientURL, advertiseClientURLPort, err := getURLAndPort(srv.Member.Etcd.AdvertiseClientURLs[0])
		if err != nil {
			return err
		}
		listenClientURL, _, err := getURLAndPort(srv.Member.Etcd.ListenClientURLs[0])
		if err != nil {
			return err
		}

		srv.advertiseClientPortToProxy[advertiseClientURLPort] = proxy.NewServer(proxy.ServerConfig{
			Logger: srv.lg,
			From:   *advertiseClientURL,
			To:     *listenClientURL,
		})
		select {
		case err = <-srv.advertiseClientPortToProxy[advertiseClientURLPort].Error():
			return err
		case <-time.After(2 * time.Second):
			srv.lg.Info("started proxy on client traffic", zap.String("url", advertiseClientURL.String()))
		}
	}

	if srv.Member.EtcdPeerProxy {
		advertisePeerURL, advertisePeerURLPort, err := getURLAndPort(srv.Member.Etcd.InitialAdvertisePeerURLs[0])
		if err != nil {
			return err
		}
		listenPeerURL, _, err := getURLAndPort(srv.Member.Etcd.ListenPeerURLs[0])
		if err != nil {
			return err
		}

		srv.advertisePeerPortToProxy[advertisePeerURLPort] = proxy.NewServer(proxy.ServerConfig{
			Logger: srv.lg,
			From:   *advertisePeerURL,
			To:     *listenPeerURL,
		})
		select {
		case err = <-srv.advertisePeerPortToProxy[advertisePeerURLPort].Error():
			return err
		case <-time.After(2 * time.Second):
			srv.lg.Info("started proxy on peer traffic", zap.String("url", advertisePeerURL.String()))
		}
	}
	return nil
}

func (srv *Server) stopProxy() {
	if srv.Member.EtcdClientProxy && len(srv.advertiseClientPortToProxy) > 0 {
		for port, px := range srv.advertiseClientPortToProxy {
			if err := px.Close(); err != nil {
				srv.lg.Warn("failed to close proxy", zap.Int("port", port))
				continue
			}
			select {
			case <-px.Done():
				// enough time to release port
				time.Sleep(time.Second)
			case <-time.After(time.Second):
			}
			srv.lg.Info("closed proxy",
				zap.Int("port", port),
				zap.String("from", px.From()),
				zap.String("to", px.To()),
			)
		}
		srv.advertiseClientPortToProxy = make(map[int]proxy.Server)
	}
	if srv.Member.EtcdPeerProxy && len(srv.advertisePeerPortToProxy) > 0 {
		for port, px := range srv.advertisePeerPortToProxy {
			if err := px.Close(); err != nil {
				srv.lg.Warn("failed to close proxy", zap.Int("port", port))
				continue
			}
			select {
			case <-px.Done():
				// enough time to release port
				time.Sleep(time.Second)
			case <-time.After(time.Second):
			}
			srv.lg.Info("closed proxy",
				zap.Int("port", port),
				zap.String("from", px.From()),
				zap.String("to", px.To()),
			)
		}
		srv.advertisePeerPortToProxy = make(map[int]proxy.Server)
	}
}

func (srv *Server) createEtcdFile() error {
	var err error
	srv.etcdLogFile, err = os.Create(srv.Member.EtcdLogPath)
	if err != nil {
		return err
	}
	srv.lg.Info("created etcd log file", zap.String("path", srv.Member.EtcdLogPath))
	return nil
}

func (srv *Server) creatEtcdCmd() {
	etcdPath, etcdFlags := srv.Member.EtcdExecPath, srv.Member.Etcd.Flags()
	u, _ := url.Parse(srv.Member.FailpointHTTPAddr)
	srv.lg.Info("creating etcd command",
		zap.String("etcd-exec-path", etcdPath),
		zap.Strings("etcd-flags", etcdFlags),
		zap.String("failpoint-http-addr", srv.Member.FailpointHTTPAddr),
		zap.String("failpoint-addr", u.Host),
	)
	srv.etcdCmd = exec.Command(etcdPath, etcdFlags...)
	srv.etcdCmd.Env = []string{"GOFAIL_HTTP=" + u.Host}
	srv.etcdCmd.Stdout = srv.etcdLogFile
	srv.etcdCmd.Stderr = srv.etcdLogFile
}

// start but do not wait for it to complete
func (srv *Server) startEtcdCmd() error {
	return srv.etcdCmd.Start()
}

func (srv *Server) handleRestartEtcd() (*rpcpb.Response, error) {
	srv.creatEtcdCmd()

	srv.lg.Info("restarting etcd")
	err := srv.startEtcdCmd()
	if err != nil {
		return nil, err
	}
	srv.lg.Info("restarted etcd", zap.String("command-path", srv.etcdCmd.Path))

	// wait some time for etcd listener start
	// before setting up proxy
	// TODO: local tests should handle port conflicts
	// with clients on restart
	time.Sleep(time.Second)
	if err = srv.startProxy(); err != nil {
		return nil, err
	}

	return &rpcpb.Response{
		Success: true,
		Status:  "successfully restarted etcd!",
	}, nil
}

func (srv *Server) handleKillEtcd() (*rpcpb.Response, error) {
	srv.stopProxy()

	err := stopWithSig(srv.etcdCmd, syscall.SIGTERM)
	if err != nil {
		return nil, err
	}
	srv.lg.Info("killed etcd", zap.String("signal", syscall.SIGTERM.String()))

	return &rpcpb.Response{
		Success: true,
		Status:  "successfully killed etcd!",
	}, nil
}

func (srv *Server) handleFailArchive() (*rpcpb.Response, error) {
	srv.stopProxy()

	// exit with stackstrace
	err := stopWithSig(srv.etcdCmd, syscall.SIGQUIT)
	if err != nil {
		return nil, err
	}
	srv.lg.Info("killed etcd", zap.String("signal", syscall.SIGQUIT.String()))

	srv.etcdLogFile.Sync()
	srv.etcdLogFile.Close()

	// TODO: support separate WAL directory
	if err = archive(
		srv.Member.BaseDir,
		srv.Member.EtcdLogPath,
		srv.Member.Etcd.DataDir,
	); err != nil {
		return nil, err
	}
	srv.lg.Info("archived data", zap.String("base-dir", srv.Member.BaseDir))

	if err = srv.createEtcdFile(); err != nil {
		return nil, err
	}

	srv.lg.Info("cleaning up page cache")
	if err := cleanPageCache(); err != nil {
		srv.lg.Warn("failed to clean up page cache", zap.String("error", err.Error()))
	}
	srv.lg.Info("cleaned up page cache")

	return &rpcpb.Response{
		Success: true,
		Status:  "successfully cleaned up etcd!",
	}, nil
}

// stop proxy, etcd, delete data directory
func (srv *Server) handleDestroyEtcdAgent() (*rpcpb.Response, error) {
	err := stopWithSig(srv.etcdCmd, syscall.SIGTERM)
	if err != nil {
		return nil, err
	}
	srv.lg.Info("killed etcd", zap.String("signal", syscall.SIGTERM.String()))

	err = os.RemoveAll(srv.Member.BaseDir)
	if err != nil {
		return nil, err
	}
	srv.lg.Info("removed base directory", zap.String("dir", srv.Member.BaseDir))

	// stop agent server
	srv.Stop()

	for port, px := range srv.advertiseClientPortToProxy {
		err := px.Close()
		srv.lg.Info("closed proxy", zap.Int("client-port", port), zap.Error(err))
	}
	for port, px := range srv.advertisePeerPortToProxy {
		err := px.Close()
		srv.lg.Info("closed proxy", zap.Int("peer-port", port), zap.Error(err))
	}

	return &rpcpb.Response{
		Success: true,
		Status:  "successfully destroyed etcd and agent!",
	}, nil
}

func (srv *Server) handleBlackholePeerPortTxRx() (*rpcpb.Response, error) {
	for port, px := range srv.advertisePeerPortToProxy {
		srv.lg.Info("blackholing", zap.Int("peer-port", port))
		px.BlackholeTx()
		px.BlackholeRx()
		srv.lg.Info("blackholed", zap.Int("peer-port", port))
	}
	return &rpcpb.Response{
		Success: true,
		Status:  "successfully blackholed peer port tx/rx!",
	}, nil
}

func (srv *Server) handleUnblackholePeerPortTxRx() (*rpcpb.Response, error) {
	for port, px := range srv.advertisePeerPortToProxy {
		srv.lg.Info("unblackholing", zap.Int("peer-port", port))
		px.UnblackholeTx()
		px.UnblackholeRx()
		srv.lg.Info("unblackholed", zap.Int("peer-port", port))
	}
	return &rpcpb.Response{
		Success: true,
		Status:  "successfully unblackholed peer port tx/rx!",
	}, nil
}

func (srv *Server) handleDelayPeerPortTxRx() (*rpcpb.Response, error) {
	lat := time.Duration(srv.Tester.UpdatedDelayLatencyMs) * time.Millisecond
	rv := time.Duration(srv.Tester.DelayLatencyMsRv) * time.Millisecond

	for port, px := range srv.advertisePeerPortToProxy {
		srv.lg.Info("delaying",
			zap.Int("peer-port", port),
			zap.Duration("latency", lat),
			zap.Duration("random-variable", rv),
		)
		px.DelayTx(lat, rv)
		px.DelayRx(lat, rv)
		srv.lg.Info("delayed",
			zap.Int("peer-port", port),
			zap.Duration("latency", lat),
			zap.Duration("random-variable", rv),
		)
	}

	return &rpcpb.Response{
		Success: true,
		Status:  "successfully delay peer port tx/rx!",
	}, nil
}

func (srv *Server) handleUndelayPeerPortTxRx() (*rpcpb.Response, error) {
	for port, px := range srv.advertisePeerPortToProxy {
		srv.lg.Info("undelaying", zap.Int("peer-port", port))
		px.UndelayTx()
		px.UndelayRx()
		srv.lg.Info("undelayed", zap.Int("peer-port", port))
	}
	return &rpcpb.Response{
		Success: true,
		Status:  "successfully undelay peer port tx/rx!",
	}, nil
}
<commit_msg>functional-tester/agent: use "AdvertisePeerURLs"<commit_after>// Copyright 2018 The etcd Authors
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package agent

import (
	"errors"
	"fmt"
	"net/url"
	"os"
	"os/exec"
	"syscall"
	"time"

	"github.com/coreos/etcd/pkg/fileutil"
	"github.com/coreos/etcd/pkg/proxy"
	"github.com/coreos/etcd/tools/functional-tester/rpcpb"

	"go.uber.org/zap"
)

// return error for system errors (e.g. fail to create files)
// return status error in response for wrong configuration/operation (e.g. start etcd twice)
func (srv *Server) handleTesterRequest(req *rpcpb.Request) (resp *rpcpb.Response, err error) {
	defer func() {
		if err == nil {
			srv.last = req.Operation
			srv.lg.Info("handler success", zap.String("operation", req.Operation.String()))
		}
	}()

	switch req.Operation {
	case rpcpb.Operation_InitialStartEtcd:
		return srv.handleInitialStartEtcd(req)
	case rpcpb.Operation_RestartEtcd:
		return srv.handleRestartEtcd()
	case rpcpb.Operation_KillEtcd:
		return srv.handleKillEtcd()
	case rpcpb.Operation_FailArchive:
		return srv.handleFailArchive()
	case rpcpb.Operation_DestroyEtcdAgent:
		return srv.handleDestroyEtcdAgent()

	case rpcpb.Operation_BlackholePeerPortTxRx:
		return srv.handleBlackholePeerPortTxRx()
	case rpcpb.Operation_UnblackholePeerPortTxRx:
		return srv.handleUnblackholePeerPortTxRx()
	case rpcpb.Operation_DelayPeerPortTxRx:
		return srv.handleDelayPeerPortTxRx()
	case rpcpb.Operation_UndelayPeerPortTxRx:
		return srv.handleUndelayPeerPortTxRx()

	default:
		msg := fmt.Sprintf("operation not found (%v)", req.Operation)
		return &rpcpb.Response{Success: false, Status: msg}, errors.New(msg)
	}
}

func (srv *Server) handleInitialStartEtcd(req *rpcpb.Request) (*rpcpb.Response, error) {
	if srv.last != rpcpb.Operation_NotStarted {
		return &rpcpb.Response{
			Success: false,
			Status:  fmt.Sprintf("%q is not valid; last server operation was %q", rpcpb.Operation_InitialStartEtcd.String(), srv.last.String()),
		}, nil
	}

	srv.Member = req.Member
	srv.Tester = req.Tester

	err := fileutil.TouchDirAll(srv.Member.BaseDir)
	if err != nil {
		return nil, err
	}
	srv.lg.Info("created base directory", zap.String("path", srv.Member.BaseDir))

	if err = srv.createEtcdFile(); err != nil {
		return nil, err
	}
	srv.creatEtcdCmd()

	err = srv.startEtcdCmd()
	if err != nil {
		return nil, err
	}
	srv.lg.Info("started etcd", zap.String("command-path", srv.etcdCmd.Path))

	// wait some time for etcd listener start
	// before setting up proxy
	time.Sleep(time.Second)
	if err = srv.startProxy(); err != nil {
		return nil, err
	}

	return &rpcpb.Response{
		Success: true,
		Status:  "successfully started etcd!",
	}, nil
}

func (srv *Server) startProxy() error {
	if srv.Member.EtcdClientProxy {
		advertiseClientURL, advertiseClientURLPort, err := getURLAndPort(srv.Member.Etcd.AdvertiseClientURLs[0])
		if err != nil {
			return err
		}
		listenClientURL, _, err := getURLAndPort(srv.Member.Etcd.ListenClientURLs[0])
		if err != nil {
			return err
		}

		srv.advertiseClientPortToProxy[advertiseClientURLPort] = proxy.NewServer(proxy.ServerConfig{
			Logger: srv.lg,
			From:   *advertiseClientURL,
			To:     *listenClientURL,
		})
		select {
		case err = <-srv.advertiseClientPortToProxy[advertiseClientURLPort].Error():
			return err
		case <-time.After(2 * time.Second):
			srv.lg.Info("started proxy on client traffic", zap.String("url", advertiseClientURL.String()))
		}
	}

	if srv.Member.EtcdPeerProxy {
		advertisePeerURL, advertisePeerURLPort, err := getURLAndPort(srv.Member.Etcd.AdvertisePeerURLs[0])
		if err != nil {
			return err
		}
		listenPeerURL, _, err := getURLAndPort(srv.Member.Etcd.ListenPeerURLs[0])
		if err != nil {
			return err
		}

		srv.advertisePeerPortToProxy[advertisePeerURLPort] = proxy.NewServer(proxy.ServerConfig{
			Logger: srv.lg,
			From:   *advertisePeerURL,
			To:     *listenPeerURL,
		})
		select {
		case err = <-srv.advertisePeerPortToProxy[advertisePeerURLPort].Error():
			return err
		case <-time.After(2 * time.Second):
			srv.lg.Info("started proxy on peer traffic", zap.String("url", advertisePeerURL.String()))
		}
	}
	return nil
}

func (srv *Server) stopProxy() {
	if srv.Member.EtcdClientProxy && len(srv.advertiseClientPortToProxy) > 0 {
		for port, px := range srv.advertiseClientPortToProxy {
			if err := px.Close(); err != nil {
				srv.lg.Warn("failed to close proxy", zap.Int("port", port))
				continue
			}
			select {
			case <-px.Done():
				// enough time to release port
				time.Sleep(time.Second)
			case <-time.After(time.Second):
			}
			srv.lg.Info("closed proxy",
				zap.Int("port", port),
				zap.String("from", px.From()),
				zap.String("to", px.To()),
			)
		}
		srv.advertiseClientPortToProxy = make(map[int]proxy.Server)
	}
	if srv.Member.EtcdPeerProxy && len(srv.advertisePeerPortToProxy) > 0 {
		for port, px := range srv.advertisePeerPortToProxy {
			if err := px.Close(); err != nil {
				srv.lg.Warn("failed to close proxy", zap.Int("port", port))
				continue
			}
			select {
			case <-px.Done():
				// enough time to release port
				time.Sleep(time.Second)
			case <-time.After(time.Second):
			}
			srv.lg.Info("closed proxy",
				zap.Int("port", port),
				zap.String("from", px.From()),
				zap.String("to", px.To()),
			)
		}
		srv.advertisePeerPortToProxy = make(map[int]proxy.Server)
	}
}

func (srv *Server) createEtcdFile() error {
	var err error
	srv.etcdLogFile, err = os.Create(srv.Member.EtcdLogPath)
	if err != nil {
		return err
	}
	srv.lg.Info("created etcd log file", zap.String("path", srv.Member.EtcdLogPath))
	return nil
}

func (srv *Server) creatEtcdCmd() {
	etcdPath, etcdFlags := srv.Member.EtcdExecPath, srv.Member.Etcd.Flags()
	u, _ := url.Parse(srv.Member.FailpointHTTPAddr)
	srv.lg.Info("creating etcd command",
		zap.String("etcd-exec-path", etcdPath),
		zap.Strings("etcd-flags", etcdFlags),
		zap.String("failpoint-http-addr", srv.Member.FailpointHTTPAddr),
		zap.String("failpoint-addr", u.Host),
	)
	srv.etcdCmd = exec.Command(etcdPath, etcdFlags...)
	srv.etcdCmd.Env = []string{"GOFAIL_HTTP=" + u.Host}
	srv.etcdCmd.Stdout = srv.etcdLogFile
	srv.etcdCmd.Stderr = srv.etcdLogFile
}

// start but do not wait for it to complete
func (srv *Server) startEtcdCmd() error {
	return srv.etcdCmd.Start()
}

func (srv *Server) handleRestartEtcd() (*rpcpb.Response, error) {
	srv.creatEtcdCmd()

	srv.lg.Info("restarting etcd")
	err := srv.startEtcdCmd()
	if err != nil {
		return nil, err
	}
	srv.lg.Info("restarted etcd", zap.String("command-path", srv.etcdCmd.Path))

	// wait some time for etcd listener start
	// before setting up proxy
	// TODO: local tests should handle port conflicts
	// with clients on restart
	time.Sleep(time.Second)
	if err = srv.startProxy(); err != nil {
		return nil, err
	}

	return &rpcpb.Response{
		Success: true,
		Status:  "successfully restarted etcd!",
	}, nil
}

func (srv *Server) handleKillEtcd() (*rpcpb.Response, error) {
	srv.stopProxy()

	err := stopWithSig(srv.etcdCmd, syscall.SIGTERM)
	if err != nil {
		return nil, err
	}
	srv.lg.Info("killed etcd", zap.String("signal", syscall.SIGTERM.String()))

	return &rpcpb.Response{
		Success: true,
		Status:  "successfully killed etcd!",
	}, nil
}

func (srv *Server) handleFailArchive() (*rpcpb.Response, error) {
	srv.stopProxy()

	// exit with stackstrace
	err := stopWithSig(srv.etcdCmd, syscall.SIGQUIT)
	if err != nil {
		return nil, err
	}
	srv.lg.Info("killed etcd", zap.String("signal", syscall.SIGQUIT.String()))

	srv.etcdLogFile.Sync()
	srv.etcdLogFile.Close()

	// TODO: support separate WAL directory
	if err = archive(
		srv.Member.BaseDir,
		srv.Member.EtcdLogPath,
		srv.Member.Etcd.DataDir,
	); err != nil {
		return nil, err
	}
	srv.lg.Info("archived data", zap.String("base-dir", srv.Member.BaseDir))

	if err = srv.createEtcdFile(); err != nil {
		return nil, err
	}

	srv.lg.Info("cleaning up page cache")
	if err := cleanPageCache(); err != nil {
		srv.lg.Warn("failed to clean up page cache", zap.String("error", err.Error()))
	}
	srv.lg.Info("cleaned up page cache")

	return &rpcpb.Response{
		Success: true,
		Status:  "successfully cleaned up etcd!",
	}, nil
}

// stop proxy, etcd, delete data directory
func (srv *Server) handleDestroyEtcdAgent() (*rpcpb.Response, error) {
	err := stopWithSig(srv.etcdCmd, syscall.SIGTERM)
	if err != nil {
		return nil, err
	}
	srv.lg.Info("killed etcd", zap.String("signal", syscall.SIGTERM.String()))

	err = os.RemoveAll(srv.Member.BaseDir)
	if err != nil {
		return nil, err
	}
	srv.lg.Info("removed base directory", zap.String("dir", srv.Member.BaseDir))

	// stop agent server
	srv.Stop()

	for port, px := range srv.advertiseClientPortToProxy {
		err := px.Close()
		srv.lg.Info("closed proxy", zap.Int("client-port", port), zap.Error(err))
	}
	for port, px := range srv.advertisePeerPortToProxy {
		err := px.Close()
		srv.lg.Info("closed proxy", zap.Int("peer-port", port), zap.Error(err))
	}

	return &rpcpb.Response{
		Success: true,
		Status:  "successfully destroyed etcd and agent!",
	}, nil
}

func (srv *Server) handleBlackholePeerPortTxRx() (*rpcpb.Response, error) {
	for port, px := range srv.advertisePeerPortToProxy {
		srv.lg.Info("blackholing", zap.Int("peer-port", port))
		px.BlackholeTx()
		px.BlackholeRx()
		srv.lg.Info("blackholed", zap.Int("peer-port", port))
	}
	return &rpcpb.Response{
		Success: true,
		Status:  "successfully blackholed peer port tx/rx!",
	}, nil
}

func (srv *Server) handleUnblackholePeerPortTxRx() (*rpcpb.Response, error) {
	for port, px := range srv.advertisePeerPortToProxy {
		srv.lg.Info("unblackholing", zap.Int("peer-port", port))
		px.UnblackholeTx()
		px.UnblackholeRx()
		srv.lg.Info("unblackholed", zap.Int("peer-port", port))
	}
	return &rpcpb.Response{
		Success: true,
		Status:  "successfully unblackholed peer port tx/rx!",
	}, nil
}

func (srv *Server) handleDelayPeerPortTxRx() (*rpcpb.Response, error) {
	lat := time.Duration(srv.Tester.UpdatedDelayLatencyMs) * time.Millisecond
	rv := time.Duration(srv.Tester.DelayLatencyMsRv) * time.Millisecond

	for port, px := range srv.advertisePeerPortToProxy {
		srv.lg.Info("delaying",
			zap.Int("peer-port", port),
			zap.Duration("latency", lat),
			zap.Duration("random-variable", rv),
		)
		px.DelayTx(lat, rv)
		px.DelayRx(lat, rv)
		srv.lg.Info("delayed",
			zap.Int("peer-port", port),
			zap.Duration("latency", lat),
			zap.Duration("random-variable", rv),
		)
	}

	return &rpcpb.Response{
		Success: true,
		Status:  "successfully delay peer port tx/rx!",
	}, nil
}

func (srv *Server) handleUndelayPeerPortTxRx() (*rpcpb.Response, error) {
	for port, px := range srv.advertisePeerPortToProxy {
		srv.lg.Info("undelaying", zap.Int("peer-port", port))
		px.UndelayTx()
		px.UndelayRx()
		srv.lg.Info("undelayed", zap.Int("peer-port", port))
	}
	return &rpcpb.Response{
		Success: true,
		Status:  "successfully undelay peer port tx/rx!",
	}, nil
}
<|endoftext|>


<commit_before>package main

import (
	"context"
	"fmt"
	"os"
	"path/filepath"
	"strings"

	"github.com/lxc/lxd/lxd/db"
	"github.com/lxc/lxd/lxd/node"
	"github.com/lxc/lxd/lxd/project"
	"github.com/lxc/lxd/lxd/rsync"
	"github.com/lxc/lxd/lxd/state"
	storagePools "github.com/lxc/lxd/lxd/storage"
	storageDrivers "github.com/lxc/lxd/lxd/storage/drivers"
	"github.com/lxc/lxd/shared"
)

func daemonStorageVolumesUnmount(s *state.State) error {
	var storageBackups string
	var storageImages string

	err := s.DB.Node.Transaction(context.TODO(), func(ctx context.Context, tx *db.NodeTx) error {
		nodeConfig, err := node.ConfigLoad(ctx, tx)
		if err != nil {
			return err
		}

		storageBackups = nodeConfig.StorageBackupsVolume()
		storageImages = nodeConfig.StorageImagesVolume()

		return nil
	})
	if err != nil {
		return err
	}

	unmount := func(storageType string, source string) error {
		// Parse the source.
		poolName, volumeName, err := daemonStorageSplitVolume(source)
		if err != nil {
			return err
		}

		pool, err := storagePools.LoadByName(s, poolName)
		if err != nil {
			return err
		}

		// Mount volume.
		_, err = pool.UnmountCustomVolume(project.Default, volumeName, nil)
		if err != nil {
			return fmt.Errorf("Failed to unmount storage volume %q: %w", source, err)
		}

		return nil
	}

	if storageBackups != "" {
		err := unmount("backups", storageBackups)
		if err != nil {
			return fmt.Errorf("Failed to unmount backups storage: %w", err)
		}
	}

	if storageImages != "" {
		err := unmount("images", storageImages)
		if err != nil {
			return fmt.Errorf("Failed to unmount images storage: %w", err)
		}
	}

	return nil
}

func daemonStorageMount(s *state.State) error {
	var storageBackups string
	var storageImages string
	err := s.DB.Node.Transaction(context.TODO(), func(ctx context.Context, tx *db.NodeTx) error {
		nodeConfig, err := node.ConfigLoad(ctx, tx)
		if err != nil {
			return err
		}

		storageBackups = nodeConfig.StorageBackupsVolume()
		storageImages = nodeConfig.StorageImagesVolume()

		return nil
	})
	if err != nil {
		return err
	}

	mount := func(storageType string, source string) error {
		// Parse the source.
		poolName, volumeName, err := daemonStorageSplitVolume(source)
		if err != nil {
			return err
		}

		pool, err := storagePools.LoadByName(s, poolName)
		if err != nil {
			return err
		}

		// Mount volume.
		err = pool.MountCustomVolume(project.Default, volumeName, nil)
		if err != nil {
			return fmt.Errorf("Failed to mount storage volume %q: %w", source, err)
		}

		return nil
	}

	if storageBackups != "" {
		err := mount("backups", storageBackups)
		if err != nil {
			return fmt.Errorf("Failed to mount backups storage: %w", err)
		}
	}

	if storageImages != "" {
		err := mount("images", storageImages)
		if err != nil {
			return fmt.Errorf("Failed to mount images storage: %w", err)
		}
	}

	return nil
}

func daemonStorageSplitVolume(volume string) (string, string, error) {
	fields := strings.Split(volume, "/")
	if len(fields) != 2 {
		return "", "", fmt.Errorf("Invalid syntax for volume, must be <pool>/<volume>")
	}

	poolName := fields[0]
	volumeName := fields[1]

	return poolName, volumeName, nil
}

func daemonStorageValidate(s *state.State, target string) error {
	// Check syntax.
	if target == "" {
		return nil
	}

	poolName, volumeName, err := daemonStorageSplitVolume(target)
	if err != nil {
		return err
	}

	// Validate pool exists.
	poolID, _, _, err := s.DB.Cluster.GetStoragePool(poolName)
	if err != nil {
		return fmt.Errorf("Unable to load storage pool %q: %w", poolName, err)
	}

	// Confirm volume exists.
	err = s.DB.Cluster.Transaction(context.TODO(), func(ctx context.Context, tx *db.ClusterTx) error {
		_, err = tx.GetStoragePoolVolume(ctx, poolID, project.Default, db.StoragePoolVolumeTypeCustom, volumeName, true)
		return err
	})
	if err != nil {
		return fmt.Errorf("Unable to load storage volume %q: %w", target, err)
	}

	snapshots, err := s.DB.Cluster.GetLocalStoragePoolVolumeSnapshotsWithType(project.Default, volumeName, db.StoragePoolVolumeTypeCustom, poolID)
	if err != nil {
		return fmt.Errorf("Unable to load storage volume snapshots %q: %w", target, err)
	}

	if len(snapshots) != 0 {
		return fmt.Errorf("Storage volumes for use by LXD itself cannot have snapshots")
	}

	pool, err := storagePools.LoadByName(s, poolName)
	if err != nil {
		return err
	}

	// Mount volume.
	err = pool.MountCustomVolume(project.Default, volumeName, nil)
	if err != nil {
		return fmt.Errorf("Failed to mount storage volume %q: %w", target, err)
	}

	defer func() { _, _ = pool.UnmountCustomVolume(project.Default, volumeName, nil) }()

	// Validate volume is empty (ignore lost+found).
	volStorageName := project.StorageVolume(project.Default, volumeName)
	mountpoint := storageDrivers.GetVolumeMountPath(poolName, storageDrivers.VolumeTypeCustom, volStorageName)

	entries, err := os.ReadDir(mountpoint)
	if err != nil {
		return fmt.Errorf("Failed to list %q: %w", mountpoint, err)
	}

	for _, entry := range entries {
		entryName := entry.Name()

		// Don't fail on clean ext4 volumes.
		if entryName == "lost+found" {
			continue
		}

		// Don't fail on systems with snapdir=visible.
		if entryName == ".zfs" {
			continue
		}

		return fmt.Errorf("Storage volume %q isn't empty", target)
	}

	return nil
}

func daemonStorageMove(s *state.State, storageType string, target string) error {
	destPath := shared.VarPath(storageType)

	// Track down the current storage.
	var sourcePool string
	var sourceVolume string

	sourcePath, err := os.Readlink(destPath)
	if err != nil {
		sourcePath = destPath
	} else {
		fields := strings.Split(sourcePath, "/")
		sourcePool = fields[len(fields)-3]
		sourceVolume = fields[len(fields)-1]
	}

	moveContent := func(source string, target string) error {
		// Copy the content.
		_, err := rsync.LocalCopy(source, target, "", false)
		if err != nil {
			return err
		}

		// Remove the source content.
		entries, err := os.ReadDir(source)
		if err != nil {
			return err
		}

		for _, entry := range entries {
			err := os.RemoveAll(filepath.Join(source, entry.Name()))
			if err != nil {
				return err
			}
		}

		return nil
	}

	// Deal with unsetting.
	if target == "" {
		// Things already look correct.
		if sourcePath == destPath {
			return nil
		}

		// Remove the symlink.
		err = os.Remove(destPath)
		if err != nil {
			return fmt.Errorf("Failed to delete storage symlink at %q: %w", destPath, err)
		}

		// Re-create as a directory.
		err = os.MkdirAll(destPath, 0700)
		if err != nil {
			return fmt.Errorf("Failed to create directory %q: %w", destPath, err)
		}

		// Move the data across.
		err = moveContent(sourcePath, destPath)
		if err != nil {
			return fmt.Errorf("Failed to move data over to directory %q: %w", destPath, err)
		}

		pool, err := storagePools.LoadByName(s, sourcePool)
		if err != nil {
			return err
		}

		// Unmount old volume.
		projectName, sourceVolumeName := project.StorageVolumeParts(sourceVolume)
		_, err = pool.UnmountCustomVolume(projectName, sourceVolumeName, nil)
		if err != nil {
			return fmt.Errorf(`Failed to umount storage volume "%s/%s": %w`, sourcePool, sourceVolumeName, err)
		}

		return nil
	}

	// Parse the target.
	poolName, volumeName, err := daemonStorageSplitVolume(target)
	if err != nil {
		return err
	}

	pool, err := storagePools.LoadByName(s, poolName)
	if err != nil {
		return err
	}

	// Mount volume.
	err = pool.MountCustomVolume(project.Default, volumeName, nil)
	if err != nil {
		return fmt.Errorf("Failed to mount storage volume %q: %w", target, err)
	}

	// Set ownership & mode.
	volStorageName := project.StorageVolume(project.Default, volumeName)
	mountpoint := storageDrivers.GetVolumeMountPath(poolName, storageDrivers.VolumeTypeCustom, volStorageName)
	destPath = mountpoint

	err = os.Chmod(mountpoint, 0700)
	if err != nil {
		return fmt.Errorf("Failed to set permissions on %q: %w", mountpoint, err)
	}

	err = os.Chown(mountpoint, 0, 0)
	if err != nil {
		return fmt.Errorf("Failed to set ownership on %q: %w", mountpoint, err)
	}

	// Handle changes.
	if sourcePath != shared.VarPath(storageType) {
		// Remove the symlink.
		err := os.Remove(shared.VarPath(storageType))
		if err != nil {
			return fmt.Errorf("Failed to remove the new symlink at %q: %w", shared.VarPath(storageType), err)
		}

		// Create the new symlink.
		err = os.Symlink(destPath, shared.VarPath(storageType))
		if err != nil {
			return fmt.Errorf("Failed to create the new symlink at %q: %w", shared.VarPath(storageType), err)
		}

		// Move the data across.
		err = moveContent(sourcePath, destPath)
		if err != nil {
			return fmt.Errorf("Failed to move data over to directory %q: %w", destPath, err)
		}

		pool, err := storagePools.LoadByName(s, sourcePool)
		if err != nil {
			return err
		}

		// Unmount old volume.
		projectName, sourceVolumeName := project.StorageVolumeParts(sourceVolume)
		_, err = pool.UnmountCustomVolume(projectName, sourceVolumeName, nil)
		if err != nil {
			return fmt.Errorf(`Failed to umount storage volume "%s/%s": %w`, sourcePool, sourceVolumeName, err)
		}

		return nil
	}

	sourcePath = shared.VarPath(storageType) + ".temp"

	// Rename the existing storage.
	err = os.Rename(shared.VarPath(storageType), sourcePath)
	if err != nil {
		return fmt.Errorf("Failed to rename existing storage %q: %w", shared.VarPath(storageType), err)
	}

	// Create the new symlink.
	err = os.Symlink(destPath, shared.VarPath(storageType))
	if err != nil {
		return fmt.Errorf("Failed to create the new symlink at %q: %w", shared.VarPath(storageType), err)
	}

	// Move the data across.
	err = moveContent(sourcePath, destPath)
	if err != nil {
		return fmt.Errorf("Failed to move data over to directory %q: %w", destPath, err)
	}

	// Remove the old data.
	err = os.RemoveAll(sourcePath)
	if err != nil {
		return fmt.Errorf("Failed to cleanup old directory %q: %w", sourcePath, err)
	}

	return nil
}
<commit_msg>lxd/daemon/storage: Improve errors in daemonStorageValidate<commit_after>package main

import (
	"context"
	"fmt"
	"os"
	"path/filepath"
	"strings"

	"github.com/lxc/lxd/lxd/db"
	"github.com/lxc/lxd/lxd/node"
	"github.com/lxc/lxd/lxd/project"
	"github.com/lxc/lxd/lxd/rsync"
	"github.com/lxc/lxd/lxd/state"
	storagePools "github.com/lxc/lxd/lxd/storage"
	storageDrivers "github.com/lxc/lxd/lxd/storage/drivers"
	"github.com/lxc/lxd/shared"
)

func daemonStorageVolumesUnmount(s *state.State) error {
	var storageBackups string
	var storageImages string

	err := s.DB.Node.Transaction(context.TODO(), func(ctx context.Context, tx *db.NodeTx) error {
		nodeConfig, err := node.ConfigLoad(ctx, tx)
		if err != nil {
			return err
		}

		storageBackups = nodeConfig.StorageBackupsVolume()
		storageImages = nodeConfig.StorageImagesVolume()

		return nil
	})
	if err != nil {
		return err
	}

	unmount := func(storageType string, source string) error {
		// Parse the source.
		poolName, volumeName, err := daemonStorageSplitVolume(source)
		if err != nil {
			return err
		}

		pool, err := storagePools.LoadByName(s, poolName)
		if err != nil {
			return err
		}

		// Mount volume.
		_, err = pool.UnmountCustomVolume(project.Default, volumeName, nil)
		if err != nil {
			return fmt.Errorf("Failed to unmount storage volume %q: %w", source, err)
		}

		return nil
	}

	if storageBackups != "" {
		err := unmount("backups", storageBackups)
		if err != nil {
			return fmt.Errorf("Failed to unmount backups storage: %w", err)
		}
	}

	if storageImages != "" {
		err := unmount("images", storageImages)
		if err != nil {
			return fmt.Errorf("Failed to unmount images storage: %w", err)
		}
	}

	return nil
}

func daemonStorageMount(s *state.State) error {
	var storageBackups string
	var storageImages string
	err := s.DB.Node.Transaction(context.TODO(), func(ctx context.Context, tx *db.NodeTx) error {
		nodeConfig, err := node.ConfigLoad(ctx, tx)
		if err != nil {
			return err
		}

		storageBackups = nodeConfig.StorageBackupsVolume()
		storageImages = nodeConfig.StorageImagesVolume()

		return nil
	})
	if err != nil {
		return err
	}

	mount := func(storageType string, source string) error {
		// Parse the source.
		poolName, volumeName, err := daemonStorageSplitVolume(source)
		if err != nil {
			return err
		}

		pool, err := storagePools.LoadByName(s, poolName)
		if err != nil {
			return err
		}

		// Mount volume.
		err = pool.MountCustomVolume(project.Default, volumeName, nil)
		if err != nil {
			return fmt.Errorf("Failed to mount storage volume %q: %w", source, err)
		}

		return nil
	}

	if storageBackups != "" {
		err := mount("backups", storageBackups)
		if err != nil {
			return fmt.Errorf("Failed to mount backups storage: %w", err)
		}
	}

	if storageImages != "" {
		err := mount("images", storageImages)
		if err != nil {
			return fmt.Errorf("Failed to mount images storage: %w", err)
		}
	}

	return nil
}

func daemonStorageSplitVolume(volume string) (string, string, error) {
	fields := strings.Split(volume, "/")
	if len(fields) != 2 {
		return "", "", fmt.Errorf("Invalid syntax for volume, must be <pool>/<volume>")
	}

	poolName := fields[0]
	volumeName := fields[1]

	return poolName, volumeName, nil
}

func daemonStorageValidate(s *state.State, target string) error {
	// Check syntax.
	if target == "" {
		return nil
	}

	poolName, volumeName, err := daemonStorageSplitVolume(target)
	if err != nil {
		return err
	}

	// Validate pool exists.
	poolID, _, _, err := s.DB.Cluster.GetStoragePool(poolName)
	if err != nil {
		return fmt.Errorf("Unable to load storage pool %q: %w", poolName, err)
	}

	// Confirm volume exists.
	err = s.DB.Cluster.Transaction(context.TODO(), func(ctx context.Context, tx *db.ClusterTx) error {
		dbVol, err := tx.GetStoragePoolVolume(ctx, poolID, project.Default, db.StoragePoolVolumeTypeCustom, volumeName, true)
		if err != nil {
			return fmt.Errorf("Failed loading storage volume %q in %q project: %w", target, project.Default, err)
		}

		if dbVol.ContentType != db.StoragePoolVolumeContentTypeNameFS {
			return fmt.Errorf("Storage volume %q in %q project is not filesystem content type", target, project.Default)
		}

		return nil
	})
	if err != nil {
		return err
	}

	snapshots, err := s.DB.Cluster.GetLocalStoragePoolVolumeSnapshotsWithType(project.Default, volumeName, db.StoragePoolVolumeTypeCustom, poolID)
	if err != nil {
		return fmt.Errorf("Unable to load storage volume snapshots %q in %q project: %w", target, project.Default, err)
	}

	if len(snapshots) != 0 {
		return fmt.Errorf("Storage volumes for use by LXD itself cannot have snapshots")
	}

	pool, err := storagePools.LoadByName(s, poolName)
	if err != nil {
		return err
	}

	// Mount volume.
	err = pool.MountCustomVolume(project.Default, volumeName, nil)
	if err != nil {
		return fmt.Errorf("Failed to mount storage volume %q: %w", target, err)
	}

	defer func() { _, _ = pool.UnmountCustomVolume(project.Default, volumeName, nil) }()

	// Validate volume is empty (ignore lost+found).
	volStorageName := project.StorageVolume(project.Default, volumeName)
	mountpoint := storageDrivers.GetVolumeMountPath(poolName, storageDrivers.VolumeTypeCustom, volStorageName)

	entries, err := os.ReadDir(mountpoint)
	if err != nil {
		return fmt.Errorf("Failed to list %q: %w", mountpoint, err)
	}

	for _, entry := range entries {
		entryName := entry.Name()

		// Don't fail on clean ext4 volumes.
		if entryName == "lost+found" {
			continue
		}

		// Don't fail on systems with snapdir=visible.
		if entryName == ".zfs" {
			continue
		}

		return fmt.Errorf("Storage volume %q isn't empty", target)
	}

	return nil
}

func daemonStorageMove(s *state.State, storageType string, target string) error {
	destPath := shared.VarPath(storageType)

	// Track down the current storage.
	var sourcePool string
	var sourceVolume string

	sourcePath, err := os.Readlink(destPath)
	if err != nil {
		sourcePath = destPath
	} else {
		fields := strings.Split(sourcePath, "/")
		sourcePool = fields[len(fields)-3]
		sourceVolume = fields[len(fields)-1]
	}

	moveContent := func(source string, target string) error {
		// Copy the content.
		_, err := rsync.LocalCopy(source, target, "", false)
		if err != nil {
			return err
		}

		// Remove the source content.
		entries, err := os.ReadDir(source)
		if err != nil {
			return err
		}

		for _, entry := range entries {
			err := os.RemoveAll(filepath.Join(source, entry.Name()))
			if err != nil {
				return err
			}
		}

		return nil
	}

	// Deal with unsetting.
	if target == "" {
		// Things already look correct.
		if sourcePath == destPath {
			return nil
		}

		// Remove the symlink.
		err = os.Remove(destPath)
		if err != nil {
			return fmt.Errorf("Failed to delete storage symlink at %q: %w", destPath, err)
		}

		// Re-create as a directory.
		err = os.MkdirAll(destPath, 0700)
		if err != nil {
			return fmt.Errorf("Failed to create directory %q: %w", destPath, err)
		}

		// Move the data across.
		err = moveContent(sourcePath, destPath)
		if err != nil {
			return fmt.Errorf("Failed to move data over to directory %q: %w", destPath, err)
		}

		pool, err := storagePools.LoadByName(s, sourcePool)
		if err != nil {
			return err
		}

		// Unmount old volume.
		projectName, sourceVolumeName := project.StorageVolumeParts(sourceVolume)
		_, err = pool.UnmountCustomVolume(projectName, sourceVolumeName, nil)
		if err != nil {
			return fmt.Errorf(`Failed to umount storage volume "%s/%s": %w`, sourcePool, sourceVolumeName, err)
		}

		return nil
	}

	// Parse the target.
	poolName, volumeName, err := daemonStorageSplitVolume(target)
	if err != nil {
		return err
	}

	pool, err := storagePools.LoadByName(s, poolName)
	if err != nil {
		return err
	}

	// Mount volume.
	err = pool.MountCustomVolume(project.Default, volumeName, nil)
	if err != nil {
		return fmt.Errorf("Failed to mount storage volume %q: %w", target, err)
	}

	// Set ownership & mode.
	volStorageName := project.StorageVolume(project.Default, volumeName)
	mountpoint := storageDrivers.GetVolumeMountPath(poolName, storageDrivers.VolumeTypeCustom, volStorageName)
	destPath = mountpoint

	err = os.Chmod(mountpoint, 0700)
	if err != nil {
		return fmt.Errorf("Failed to set permissions on %q: %w", mountpoint, err)
	}

	err = os.Chown(mountpoint, 0, 0)
	if err != nil {
		return fmt.Errorf("Failed to set ownership on %q: %w", mountpoint, err)
	}

	// Handle changes.
	if sourcePath != shared.VarPath(storageType) {
		// Remove the symlink.
		err := os.Remove(shared.VarPath(storageType))
		if err != nil {
			return fmt.Errorf("Failed to remove the new symlink at %q: %w", shared.VarPath(storageType), err)
		}

		// Create the new symlink.
		err = os.Symlink(destPath, shared.VarPath(storageType))
		if err != nil {
			return fmt.Errorf("Failed to create the new symlink at %q: %w", shared.VarPath(storageType), err)
		}

		// Move the data across.
		err = moveContent(sourcePath, destPath)
		if err != nil {
			return fmt.Errorf("Failed to move data over to directory %q: %w", destPath, err)
		}

		pool, err := storagePools.LoadByName(s, sourcePool)
		if err != nil {
			return err
		}

		// Unmount old volume.
		projectName, sourceVolumeName := project.StorageVolumeParts(sourceVolume)
		_, err = pool.UnmountCustomVolume(projectName, sourceVolumeName, nil)
		if err != nil {
			return fmt.Errorf(`Failed to umount storage volume "%s/%s": %w`, sourcePool, sourceVolumeName, err)
		}

		return nil
	}

	sourcePath = shared.VarPath(storageType) + ".temp"

	// Rename the existing storage.
	err = os.Rename(shared.VarPath(storageType), sourcePath)
	if err != nil {
		return fmt.Errorf("Failed to rename existing storage %q: %w", shared.VarPath(storageType), err)
	}

	// Create the new symlink.
	err = os.Symlink(destPath, shared.VarPath(storageType))
	if err != nil {
		return fmt.Errorf("Failed to create the new symlink at %q: %w", shared.VarPath(storageType), err)
	}

	// Move the data across.
	err = moveContent(sourcePath, destPath)
	if err != nil {
		return fmt.Errorf("Failed to move data over to directory %q: %w", destPath, err)
	}

	// Remove the old data.
	err = os.RemoveAll(sourcePath)
	if err != nil {
		return fmt.Errorf("Failed to cleanup old directory %q: %w", sourcePath, err)
	}

	return nil
}
<|endoftext|>


<commit_before>package dispatcher

//go:generate errorgen

import (
	"context"
	"strings"
	"sync"
	"time"

	"v2ray.com/core"
	"v2ray.com/core/app/proxyman"
	"v2ray.com/core/common"
	"v2ray.com/core/common/buf"
	"v2ray.com/core/common/net"
	"v2ray.com/core/common/protocol"
	"v2ray.com/core/common/session"
	"v2ray.com/core/features/outbound"
	"v2ray.com/core/features/policy"
	"v2ray.com/core/features/routing"
	"v2ray.com/core/features/stats"
	"v2ray.com/core/transport"
	"v2ray.com/core/transport/pipe"
)

var (
	errSniffingTimeout = newError("timeout on sniffing")
)

type cachedReader struct {
	sync.Mutex
	reader *pipe.Reader
	cache  buf.MultiBuffer
}

func (r *cachedReader) Cache(b *buf.Buffer) {
	mb, _ := r.reader.ReadMultiBufferTimeout(time.Millisecond * 100)
	r.Lock()
	if !mb.IsEmpty() {
		r.cache, _ = buf.MergeMulti(r.cache, mb)
	}
	b.Clear()
	rawBytes := b.Extend(buf.Size)
	n := r.cache.Copy(rawBytes)
	b.Resize(0, int32(n))
	r.Unlock()
}

func (r *cachedReader) readInternal() buf.MultiBuffer {
	r.Lock()
	defer r.Unlock()

	if r.cache != nil && !r.cache.IsEmpty() {
		mb := r.cache
		r.cache = nil
		return mb
	}

	return nil
}

func (r *cachedReader) ReadMultiBuffer() (buf.MultiBuffer, error) {
	mb := r.readInternal()
	if mb != nil {
		return mb, nil
	}

	return r.reader.ReadMultiBuffer()
}

func (r *cachedReader) ReadMultiBufferTimeout(timeout time.Duration) (buf.MultiBuffer, error) {
	mb := r.readInternal()
	if mb != nil {
		return mb, nil
	}

	return r.reader.ReadMultiBufferTimeout(timeout)
}

func (r *cachedReader) Interrupt() {
	r.Lock()
	if r.cache != nil {
		r.cache = buf.ReleaseMulti(r.cache)
	}
	r.Unlock()
	r.reader.Interrupt()
}

// DefaultDispatcher is a default implementation of Dispatcher.
type DefaultDispatcher struct {
	ohm    outbound.Manager
	router routing.Router
	policy policy.Manager
	stats  stats.Manager
}

func init() {
	common.Must(common.RegisterConfig((*Config)(nil), func(ctx context.Context, config interface{}) (interface{}, error) {
		d := new(DefaultDispatcher)
		if err := core.RequireFeatures(ctx, func(om outbound.Manager, router routing.Router, pm policy.Manager, sm stats.Manager) error {
			return d.Init(config.(*Config), om, router, pm, sm)
		}); err != nil {
			return nil, err
		}
		return d, nil
	}))
}

// Init initializes DefaultDispatcher.
func (d *DefaultDispatcher) Init(config *Config, om outbound.Manager, router routing.Router, pm policy.Manager, sm stats.Manager) error {
	d.ohm = om
	d.router = router
	d.policy = pm
	d.stats = sm
	return nil
}

// Type implements common.HasType.
func (*DefaultDispatcher) Type() interface{} {
	return routing.DispatcherType()
}

// Start implements common.Runnable.
func (*DefaultDispatcher) Start() error {
	return nil
}

// Close implements common.Closable.
func (*DefaultDispatcher) Close() error { return nil }

func (d *DefaultDispatcher) getLink(ctx context.Context) (*transport.Link, *transport.Link) {
	opt := pipe.OptionsFromContext(ctx)
	uplinkReader, uplinkWriter := pipe.New(opt...)
	downlinkReader, downlinkWriter := pipe.New(opt...)

	inboundLink := &transport.Link{
		Reader: downlinkReader,
		Writer: uplinkWriter,
	}

	outboundLink := &transport.Link{
		Reader: uplinkReader,
		Writer: downlinkWriter,
	}

	sessionInbound := session.InboundFromContext(ctx)
	var user *protocol.MemoryUser
	if sessionInbound != nil {
		user = sessionInbound.User
	}

	if user != nil && len(user.Email) > 0 {
		p := d.policy.ForLevel(user.Level)
		if p.Stats.UserUplink {
			name := "user>>>" + user.Email + ">>>traffic>>>uplink"
			if c, _ := stats.GetOrRegisterCounter(d.stats, name); c != nil {
				inboundLink.Writer = &SizeStatWriter{
					Counter: c,
					Writer:  inboundLink.Writer,
				}
			}
		}
		if p.Stats.UserDownlink {
			name := "user>>>" + user.Email + ">>>traffic>>>downlink"
			if c, _ := stats.GetOrRegisterCounter(d.stats, name); c != nil {
				outboundLink.Writer = &SizeStatWriter{
					Counter: c,
					Writer:  outboundLink.Writer,
				}
			}
		}
	}

	return inboundLink, outboundLink
}

func shouldOverride(result SniffResult, domainOverride []string) bool {
	for _, p := range domainOverride {
		if strings.HasPrefix(result.Protocol(), p) {
			return true
		}
	}
	return false
}

// Dispatch implements routing.Dispatcher.
func (d *DefaultDispatcher) Dispatch(ctx context.Context, destination net.Destination) (*transport.Link, error) {
	if !destination.IsValid() {
		panic("Dispatcher: Invalid destination.")
	}
	ob := &session.Outbound{
		Target: destination,
	}
	ctx = session.ContextWithOutbound(ctx, ob)

	inbound, outbound := d.getLink(ctx)
	sniffingConfig := proxyman.SniffingConfigFromContext(ctx)
	if destination.Network != net.Network_TCP || sniffingConfig == nil || !sniffingConfig.Enabled {
		go d.routedDispatch(ctx, outbound, destination)
	} else {
		go func() {
			cReader := &cachedReader{
				reader: outbound.Reader.(*pipe.Reader),
			}
			outbound.Reader = cReader
			result, err := sniffer(ctx, cReader)
			if err == nil {
				ctx = ContextWithSniffingResult(ctx, result)
			}
			if err == nil && shouldOverride(result, sniffingConfig.DestinationOverride) {
				domain := result.Domain()
				newError("sniffed domain: ", domain).WriteToLog(session.ExportIDToError(ctx))
				destination.Address = net.ParseAddress(domain)
				ob.Target = destination
			}
			d.routedDispatch(ctx, outbound, destination)
		}()
	}
	return inbound, nil
}

func sniffer(ctx context.Context, cReader *cachedReader) (SniffResult, error) {
	payload := buf.New()
	defer payload.Release()

	sniffer := NewSniffer()
	totalAttempt := 0
	for {
		select {
		case <-ctx.Done():
			return nil, ctx.Err()
		default:
			totalAttempt++
			if totalAttempt > 2 {
				return nil, errSniffingTimeout
			}

			cReader.Cache(payload)
			if !payload.IsEmpty() {
				result, err := sniffer.Sniff(payload.Bytes())
				if err != common.ErrNoClue {
					return result, err
				}
			}
			if payload.IsFull() {
				return nil, errUnknownContent
			}
		}
	}
}

func (d *DefaultDispatcher) routedDispatch(ctx context.Context, link *transport.Link, destination net.Destination) {
	dispatcher := d.ohm.GetDefaultHandler()
	if d.router != nil {
		if tag, err := d.router.PickRoute(ctx); err == nil {
			if handler := d.ohm.GetHandler(tag); handler != nil {
				newError("taking detour [", tag, "] for [", destination, "]").WriteToLog(session.ExportIDToError(ctx))
				dispatcher = handler
			} else {
				newError("non existing tag: ", tag).AtWarning().WriteToLog(session.ExportIDToError(ctx))
			}
		} else {
			newError("default route for ", destination).WriteToLog(session.ExportIDToError(ctx))
		}
	}

	if dispatcher == nil {
		newError("default outbound handler not exist").WriteToLog(session.ExportIDToError(ctx))
		common.Close(link.Writer)
		common.Interrupt(link.Reader)
		return
	}

	dispatcher.Dispatch(ctx, link)
}
<commit_msg>rename dispatcher to handler<commit_after>package dispatcher

//go:generate errorgen

import (
	"context"
	"strings"
	"sync"
	"time"

	"v2ray.com/core"
	"v2ray.com/core/app/proxyman"
	"v2ray.com/core/common"
	"v2ray.com/core/common/buf"
	"v2ray.com/core/common/net"
	"v2ray.com/core/common/protocol"
	"v2ray.com/core/common/session"
	"v2ray.com/core/features/outbound"
	"v2ray.com/core/features/policy"
	"v2ray.com/core/features/routing"
	"v2ray.com/core/features/stats"
	"v2ray.com/core/transport"
	"v2ray.com/core/transport/pipe"
)

var (
	errSniffingTimeout = newError("timeout on sniffing")
)

type cachedReader struct {
	sync.Mutex
	reader *pipe.Reader
	cache  buf.MultiBuffer
}

func (r *cachedReader) Cache(b *buf.Buffer) {
	mb, _ := r.reader.ReadMultiBufferTimeout(time.Millisecond * 100)
	r.Lock()
	if !mb.IsEmpty() {
		r.cache, _ = buf.MergeMulti(r.cache, mb)
	}
	b.Clear()
	rawBytes := b.Extend(buf.Size)
	n := r.cache.Copy(rawBytes)
	b.Resize(0, int32(n))
	r.Unlock()
}

func (r *cachedReader) readInternal() buf.MultiBuffer {
	r.Lock()
	defer r.Unlock()

	if r.cache != nil && !r.cache.IsEmpty() {
		mb := r.cache
		r.cache = nil
		return mb
	}

	return nil
}

func (r *cachedReader) ReadMultiBuffer() (buf.MultiBuffer, error) {
	mb := r.readInternal()
	if mb != nil {
		return mb, nil
	}

	return r.reader.ReadMultiBuffer()
}

func (r *cachedReader) ReadMultiBufferTimeout(timeout time.Duration) (buf.MultiBuffer, error) {
	mb := r.readInternal()
	if mb != nil {
		return mb, nil
	}

	return r.reader.ReadMultiBufferTimeout(timeout)
}

func (r *cachedReader) Interrupt() {
	r.Lock()
	if r.cache != nil {
		r.cache = buf.ReleaseMulti(r.cache)
	}
	r.Unlock()
	r.reader.Interrupt()
}

// DefaultDispatcher is a default implementation of Dispatcher.
type DefaultDispatcher struct {
	ohm    outbound.Manager
	router routing.Router
	policy policy.Manager
	stats  stats.Manager
}

func init() {
	common.Must(common.RegisterConfig((*Config)(nil), func(ctx context.Context, config interface{}) (interface{}, error) {
		d := new(DefaultDispatcher)
		if err := core.RequireFeatures(ctx, func(om outbound.Manager, router routing.Router, pm policy.Manager, sm stats.Manager) error {
			return d.Init(config.(*Config), om, router, pm, sm)
		}); err != nil {
			return nil, err
		}
		return d, nil
	}))
}

// Init initializes DefaultDispatcher.
func (d *DefaultDispatcher) Init(config *Config, om outbound.Manager, router routing.Router, pm policy.Manager, sm stats.Manager) error {
	d.ohm = om
	d.router = router
	d.policy = pm
	d.stats = sm
	return nil
}

// Type implements common.HasType.
func (*DefaultDispatcher) Type() interface{} {
	return routing.DispatcherType()
}

// Start implements common.Runnable.
func (*DefaultDispatcher) Start() error {
	return nil
}

// Close implements common.Closable.
func (*DefaultDispatcher) Close() error { return nil }

func (d *DefaultDispatcher) getLink(ctx context.Context) (*transport.Link, *transport.Link) {
	opt := pipe.OptionsFromContext(ctx)
	uplinkReader, uplinkWriter := pipe.New(opt...)
	downlinkReader, downlinkWriter := pipe.New(opt...)

	inboundLink := &transport.Link{
		Reader: downlinkReader,
		Writer: uplinkWriter,
	}

	outboundLink := &transport.Link{
		Reader: uplinkReader,
		Writer: downlinkWriter,
	}

	sessionInbound := session.InboundFromContext(ctx)
	var user *protocol.MemoryUser
	if sessionInbound != nil {
		user = sessionInbound.User
	}

	if user != nil && len(user.Email) > 0 {
		p := d.policy.ForLevel(user.Level)
		if p.Stats.UserUplink {
			name := "user>>>" + user.Email + ">>>traffic>>>uplink"
			if c, _ := stats.GetOrRegisterCounter(d.stats, name); c != nil {
				inboundLink.Writer = &SizeStatWriter{
					Counter: c,
					Writer:  inboundLink.Writer,
				}
			}
		}
		if p.Stats.UserDownlink {
			name := "user>>>" + user.Email + ">>>traffic>>>downlink"
			if c, _ := stats.GetOrRegisterCounter(d.stats, name); c != nil {
				outboundLink.Writer = &SizeStatWriter{
					Counter: c,
					Writer:  outboundLink.Writer,
				}
			}
		}
	}

	return inboundLink, outboundLink
}

func shouldOverride(result SniffResult, domainOverride []string) bool {
	for _, p := range domainOverride {
		if strings.HasPrefix(result.Protocol(), p) {
			return true
		}
	}
	return false
}

// Dispatch implements routing.Dispatcher.
func (d *DefaultDispatcher) Dispatch(ctx context.Context, destination net.Destination) (*transport.Link, error) {
	if !destination.IsValid() {
		panic("Dispatcher: Invalid destination.")
	}
	ob := &session.Outbound{
		Target: destination,
	}
	ctx = session.ContextWithOutbound(ctx, ob)

	inbound, outbound := d.getLink(ctx)
	sniffingConfig := proxyman.SniffingConfigFromContext(ctx)
	if destination.Network != net.Network_TCP || sniffingConfig == nil || !sniffingConfig.Enabled {
		go d.routedDispatch(ctx, outbound, destination)
	} else {
		go func() {
			cReader := &cachedReader{
				reader: outbound.Reader.(*pipe.Reader),
			}
			outbound.Reader = cReader
			result, err := sniffer(ctx, cReader)
			if err == nil {
				ctx = ContextWithSniffingResult(ctx, result)
			}
			if err == nil && shouldOverride(result, sniffingConfig.DestinationOverride) {
				domain := result.Domain()
				newError("sniffed domain: ", domain).WriteToLog(session.ExportIDToError(ctx))
				destination.Address = net.ParseAddress(domain)
				ob.Target = destination
			}
			d.routedDispatch(ctx, outbound, destination)
		}()
	}
	return inbound, nil
}

func sniffer(ctx context.Context, cReader *cachedReader) (SniffResult, error) {
	payload := buf.New()
	defer payload.Release()

	sniffer := NewSniffer()
	totalAttempt := 0
	for {
		select {
		case <-ctx.Done():
			return nil, ctx.Err()
		default:
			totalAttempt++
			if totalAttempt > 2 {
				return nil, errSniffingTimeout
			}

			cReader.Cache(payload)
			if !payload.IsEmpty() {
				result, err := sniffer.Sniff(payload.Bytes())
				if err != common.ErrNoClue {
					return result, err
				}
			}
			if payload.IsFull() {
				return nil, errUnknownContent
			}
		}
	}
}

func (d *DefaultDispatcher) routedDispatch(ctx context.Context, link *transport.Link, destination net.Destination) {
	var handler outbound.Handler
	if d.router != nil {
		if tag, err := d.router.PickRoute(ctx); err == nil {
			if h := d.ohm.GetHandler(tag); h != nil {
				newError("taking detour [", tag, "] for [", destination, "]").WriteToLog(session.ExportIDToError(ctx))
				handler = h
			} else {
				newError("non existing tag: ", tag).AtWarning().WriteToLog(session.ExportIDToError(ctx))
			}
		} else {
			newError("default route for ", destination).WriteToLog(session.ExportIDToError(ctx))
		}
	}

	if handler == nil {
		handler = d.ohm.GetDefaultHandler()
	}

	if handler == nil {
		newError("default outbound handler not exist").WriteToLog(session.ExportIDToError(ctx))
		common.Close(link.Writer)
		common.Interrupt(link.Reader)
		return
	}

	handler.Dispatch(ctx, link)
}
<|endoftext|>