<commit_before>//
// nntp.go
//
package srnd

import (
  "bufio"
  "io"
  "log"
  "net"
  "net/textproto"
  "os"
  "strings"
  "time"
)

type ConnectionInfo struct {
  mode string
  newsgroup string
  allowsPosting bool
  supportsStream bool
  state string
  // if true we are reading data
  reading bool
}

type NNTPConnection struct {
  conn net.Conn
  txtconn *textproto.Conn
  inbound bool
  debug bool
  info *ConnectionInfo
  policy *FeedPolicy
  // channel for senging sync messages
  sync chan string
  // message io
  msg_reader MessageReader
  msg_writer MessageWriter
  // do we allow articles from tor?
  allow_tor bool
  allow_tor_attachments bool
}

// ask if they need this article
func (self *NNTPConnection) askSync(msgid string) {
  if ValidMessageID(msgid) {
    self.txtconn.PrintfLine("CHECK %s", msgid)
  }
}

func (self *NNTPConnection) HandleOutbound(d *NNTPDaemon, quarks map[string]string) {
  code, line, _ := self.txtconn.ReadCodeLine(-1)
  self.info.allowsPosting = code == 200
  if ! self.info.allowsPosting {
    log.Printf("outbound feed posting not allowed: %d %s", code, line)
    self.Close()
    return
  }
  // TODO: autodetect quarks
  if quarks != nil {
    // check for twisted news server quark
    t, ok := quarks["twisted"]
    if ok {
      if t == "1" {
        // force into post mode
        self.info.mode = "post"
      }
    }
  }

  if self.info.mode == "" {
    // they allow posting
    // send capabilities command
    _ = self.txtconn.PrintfLine("CAPABILITIES")
    capreader := bufio.NewReader(self.txtconn.DotReader())

    // get capabilites
    for {
      line, err := capreader.ReadString('\n')
      if err != nil {
        log.Println(err)
        break
      }
      line = strings.ToLower(line)
      if line == "streaming\n" {
        self.info.supportsStream = true
      } else if line == "postihavestreaming\n" {
        self.info.supportsStream = true
      } else {
        // wtf now?
      }
    }
  }
  // if they support streaming and allow posting continue
  // otherwise quit
  if self.info.supportsStream && self.info.allowsPosting {
    self.streaming_mode(d)
  } else if self.info.mode == "post" {
    // we are forced into post mode
    self.post_mode(d)
  }
}

// enter posting mode
func (self *NNTPConnection) post_mode(d *NNTPDaemon) {
  // TODO: improve this, it's dumb right now
  for {
    msg_id := <- self.sync
    fname := d.store.GetFilename(msg_id)
    f, err := os.Open(fname)
    if f == nil {
      continue
    }
    self.txtconn.PrintfLine("POST")

    _ , _, err = self.txtconn.ReadCodeLine(340)
    if err == nil {
      w := self.txtconn.DotWriter()
      _, err = io.Copy(w, f)
      w.Close()
      _, _, err = self.txtconn.ReadCodeLine(240)
      if err != nil {
        log.Println("failed to send article",msg_id, "via POST", err)
      }
    }
    f.Close()
  }
}

// enter streaming mode
func (self *NNTPConnection) streaming_mode(d *NNTPDaemon) {
  var err error
  var line string
  var code int
  err = self.txtconn.PrintfLine("MODE STREAM")
  if err != nil {
    log.Println("failed to initiated streaming mode on feed", err)
    return
  }
  code, line, err = self.txtconn.ReadCodeLine(-1)
  if err != nil {
    log.Println("failed to read response for streaming handshake on feed", err)
    return
  }
  if code == 203 {
    self.info.mode = "stream"
    log.Println("streaming mode activated")
  } else {
    log.Println("streaming mode not activated, quitting")
    self.Quit()
    return
  }
  log.Println("outfeed enter mainloop")

  // go routine for sending sync requests
  go func() {
    for {
      msg_id := <- self.sync
      for self.info.reading {
        time.Sleep(10 * time.Millisecond)
      }
      self.askSync(msg_id)
    }
  }()

  for {
    code, line, err = self.txtconn.ReadCodeLine(-1)
    if err != nil {
      log.Println("error reading response code", err)
      return
    }
    code = int(code)
    commands := strings.Split(line, " ")
    if code == 238 && len(commands) > 1 && ValidMessageID(commands[0]) {
      fname := d.store.GetFilename(commands[0])
      f, err := os.Open(fname)
      if f == nil {
        continue
      }
      err = self.SendMessage(commands[0], f, d)
      f.Close()
      if err != nil {
        log.Println("failed to send message", err)
        self.Quit()
        return
      }
    } else if code == 438 {
      // declined
      continue
    } else if code == 239 {
      // accepted
      continue
    } else if code == 439 {
      // invalid
      continue
    } else {
      log.Printf("invalid response from outbound feed: '%d %s'", code, line)
    }
  }
}

// just do it (tm)
func (self *NNTPConnection) SendMessage(msgid string, msg io.Reader, d *NNTPDaemon) error {
  var err error
  self.info.reading = true
  err = self.txtconn.PrintfLine("TAKETHIS %s", msgid)
  if err != nil {
    log.Println("error in outfeed", err)
    return  err
  }
  wr := self.txtconn.DotWriter()
  _, err = io.Copy(wr, msg)
  wr.Close()
  self.info.reading = false
  if err != nil {
    log.Printf("failed to send %s via feed: %s", msgid, err)
    return err
  }
  return nil
}

// handle inbound connection
func (self *NNTPConnection) HandleInbound(d *NNTPDaemon) {

  // intitiate handshake
  var err error
  self.info.mode = "STREAM"
  log.Println("Incoming nntp connection from", self.conn.RemoteAddr())
  // send welcome
  greet := "2nd generation overchan NNTP Daemon posting allowed"
  self.txtconn.PrintfLine("200 %s", greet)
  for {
    if err != nil {
      log.Println("failure in infeed", err)
      self.Quit()
      return
    }
    // read line and break if needed
    line, err := self.ReadLine()
    if len(line) == 0 || err != nil {
      break
    }
    var code int
    var msg string
    commands := strings.Split(line, " ")
    cmd := commands[0]
    // capabilities command
    if cmd == "CAPABILITIES" {
      self.sendCapabilities()
    } else if cmd == "MODE" { // mode switch
      if len(commands) == 2 {
        // get mode
        mode := strings.ToUpper(commands[1])
        // reader mode
        if mode == "READER" {
          self.info.mode = "reader"
          code = 201
          msg = "posting disallowed"
        } else if mode == "STREAM" {
          // mode stream
          self.info.mode = "stream"
          code = 203
          msg = "stream it"
        } else {
          // other modes not implemented
          code = 501
          msg = "mode not implemented"
        }
      } else {
        code = 500
        msg = "syntax error"
      }

      self.txtconn.PrintfLine("%d %s", code, msg)
    } else if self.info.mode == "stream" { // we are in stream mode
      if cmd == "TAKETHIS" {
        var newsgroup string
        if len(commands) == 2 {
          article := commands[1]
          if ValidMessageID(article) {
            code := 239
            file := d.store.CreateTempFile(article)
            ip_header := ""
            ip_banned := false
            headers_done := false
            read_more := true
            has_attachment := false
            is_signed := false
            message := "we are gud"
            for {
              var line string
              line, err = self.ReadLine()
              if err != nil {
                log.Println("error reading", article, err)
                break
              }

              if line == "" && ! headers_done {
                // headers done
                headers_done = true
                if len(ip_header) > 0 {
                  ip_banned, err = d.database.CheckEncIPBanned(ip_header)
                  if err == nil {
                    if ip_banned {
                      code = 438
                      read_more = false
                      message = "poster is banned"
                    }
                  } else {
                    log.Println("cannot check for banned encrypted ip", err)
                    // send it later
                    code = 431
                    message = "could be banned by us but we do not know"
                  }
                } else if self.allow_tor {
                  // do we want tor posts with attachments?
                  if has_attachment && ! self.allow_tor_attachments {
                    // this guy is banned
                    code = 438
                    read_more = false
                    message = "we do not take attachments from tor"
                  }
                } else {
                  // we don'e want it
                  code = 438
                  message = "we do not take anonymous posts"
                  read_more = false
                }
                if is_signed {
                  log.Println("we got a signed message")
                }
              }

              if line == "." {
                break
              } else if ! headers_done {
                lower_line := strings.ToLower(line)
                // newsgroup header
                if strings.HasPrefix(lower_line, "newsgroups: ") {
                  if len(newsgroup) == 0 {
                    newsgroup := line[12:]
                    if ! newsgroupValidFormat(newsgroup) {
                      // bad newsgroup
                      code = 438
                      message = "invalid newsgroup"
                    }
                  }
                } else if strings.HasPrefix(lower_line, "x-tor-poster: 1") {
                  if ! self.allow_tor {
                    // we don't want this post
                    code = 438
                    message = "we do not take anonymous posts"
                    read_more = false
                  }
                } else if strings.HasPrefix(lower_line, "x-encrypted-ip: ") {
                  ip_header = strings.Split(line, " ")[1]
                  ip_header = strings.Trim(line, " \t\r\n")
                } else if strings.HasPrefix(lower_line, "content-type: multipart") {
                  has_attachment = true
                } else if strings.HasPrefix(lower_line, "x-signature-ed25519-sha512: ") {
                  is_signed = true
                  has_attachment = true
                }
              }
              if read_more {
                file.Write([]byte(line))
                file.Write([]byte("\n"))
              }
            }
            file.Close()
            // tell them our result
            self.txtconn.PrintfLine("%d %s %s", code, article, message)
            // the send was good
            if code == 239 {
              log.Println(self.conn.RemoteAddr(), "got article", article)
              // inform daemon
              d.infeed_load <- article
            } else if code == 431 {
              // defer
              log.Println("defering receival, keeping article in tmp dir:", article)
            } else {
              // delete unaccepted article
              log.Println("did not accept", article)
              fname := d.store.GetTempFilename(article)
              DelFile(fname)
            }
          } else {
            self.txtconn.PrintfLine("438 %s invalid message-id", article)
          }
        }
      }
      // check command
      if cmd == "CHECK" {
        if len(commands) == 2 {
          // check syntax
          // send error if needed
          article := commands[1]
          if ! ValidMessageID(article) {
            self.txtconn.PrintfLine("501 bad message id")
            continue
          }
          // do we already have this article?
          if d.store.HasArticle(article) {
            // ya, we got it already
            // tell them to not send it
            self.txtconn.PrintfLine("438 %s we have this article", article)
          } else {
            // nope, we do not have it
            // tell them to send it
            self.txtconn.PrintfLine("238 %s we want this article please give it", article)
          }
        }
      }
    }
  }
  self.Close()
}

func (self *NNTPConnection) sendCapabilities() {
  wr := self.txtconn.DotWriter()
  io.WriteString(wr, "101 we can haz do things\r\n")
  io.WriteString(wr, "VERSION 2\r\n")
  io.WriteString(wr, "IMPLEMENTATION srndv2 better than SRNd\r\n")
  io.WriteString(wr, "STREAMING\r\n")
  io.WriteString(wr, "READER\r\n")
  wr.Close()
}

func (self *NNTPConnection) Quit() {
  if ! self.inbound {
    self.txtconn.PrintfLine("QUIT")
  }
  self.Close()
}

func (self *NNTPConnection) ReadLine() (string, error) {
  line, err := self.txtconn.ReadLine()
  if err != nil {
    log.Println("error reading line in feed", err)
    return "", err
  }
  if self.debug {
    log.Println(self.conn.RemoteAddr(), "recv line", line)
  }
  return line, nil
}

// close the connection
func (self *NNTPConnection) Close() {
  err := self.conn.Close()
  if err != nil {
    log.Println(self.conn.RemoteAddr(), err)
  }
  log.Println(self.conn.RemoteAddr(), "Closed Connection")
}
<commit_msg>try adding more logging to post method<commit_after>//
// nntp.go
//
package srnd

import (
  "bufio"
  "io"
  "log"
  "net"
  "net/textproto"
  "os"
  "strings"
  "time"
)

type ConnectionInfo struct {
  mode string
  newsgroup string
  allowsPosting bool
  supportsStream bool
  state string
  // if true we are reading data
  reading bool
}

type NNTPConnection struct {
  conn net.Conn
  txtconn *textproto.Conn
  inbound bool
  debug bool
  info *ConnectionInfo
  policy *FeedPolicy
  // channel for senging sync messages
  sync chan string
  // message io
  msg_reader MessageReader
  msg_writer MessageWriter
  // do we allow articles from tor?
  allow_tor bool
  allow_tor_attachments bool
}

// ask if they need this article
func (self *NNTPConnection) askSync(msgid string) {
  if ValidMessageID(msgid) {
    self.txtconn.PrintfLine("CHECK %s", msgid)
  }
}

func (self *NNTPConnection) HandleOutbound(d *NNTPDaemon, quarks map[string]string) {
  code, line, _ := self.txtconn.ReadCodeLine(-1)
  self.info.allowsPosting = code == 200
  if ! self.info.allowsPosting {
    log.Printf("outbound feed posting not allowed: %d %s", code, line)
    self.Close()
    return
  }
  // TODO: autodetect quarks
  if quarks != nil {
    // check for twisted news server quark
    t, ok := quarks["twisted"]
    if ok {
      if t == "1" {
        // force into post mode
        self.info.mode = "post"
      }
    }
  }

  if self.info.mode == "" {
    // they allow posting
    // send capabilities command
    _ = self.txtconn.PrintfLine("CAPABILITIES")
    capreader := bufio.NewReader(self.txtconn.DotReader())

    // get capabilites
    for {
      line, err := capreader.ReadString('\n')
      if err != nil {
        log.Println(err)
        break
      }
      line = strings.ToLower(line)
      if line == "streaming\n" {
        self.info.supportsStream = true
      } else if line == "postihavestreaming\n" {
        self.info.supportsStream = true
      } else {
        // wtf now?
      }
    }
  }
  // if they support streaming and allow posting continue
  // otherwise quit
  if self.info.supportsStream && self.info.allowsPosting {
    self.streaming_mode(d)
  } else if self.info.mode == "post" {
    // we are forced into post mode
    self.post_mode(d)
  }
}

// enter posting mode
func (self *NNTPConnection) post_mode(d *NNTPDaemon) {
  // TODO: improve this, it's dumb right now
  for {
    msg_id := <- self.sync
    fname := d.store.GetFilename(msg_id)
    f, err := os.Open(fname)
    if f == nil {
      continue
    }
    self.txtconn.PrintfLine("POST")

    _ , _, err = self.txtconn.ReadCodeLine(340)
    if err == nil {
      w := self.txtconn.DotWriter()
      _, err = io.Copy(w, f)
      w.Close()
      _, _, err = self.txtconn.ReadCodeLine(240)
      if err != nil {
        log.Println("failed to send article",msg_id, "via POST", err)
        self.Quit()
        return
      }
    }
    f.Close()
  }
}

// enter streaming mode
func (self *NNTPConnection) streaming_mode(d *NNTPDaemon) {
  var err error
  var line string
  var code int
  err = self.txtconn.PrintfLine("MODE STREAM")
  if err != nil {
    log.Println("failed to initiated streaming mode on feed", err)
    return
  }
  code, line, err = self.txtconn.ReadCodeLine(-1)
  if err != nil {
    log.Println("failed to read response for streaming handshake on feed", err)
    return
  }
  if code == 203 {
    self.info.mode = "stream"
    log.Println("streaming mode activated")
  } else {
    log.Println("streaming mode not activated, quitting")
    self.Quit()
    return
  }
  log.Println("outfeed enter mainloop")

  // go routine for sending sync requests
  go func() {
    for {
      msg_id := <- self.sync
      for self.info.reading {
        time.Sleep(10 * time.Millisecond)
      }
      self.askSync(msg_id)
    }
  }()

  for {
    code, line, err = self.txtconn.ReadCodeLine(-1)
    if err != nil {
      log.Println("error reading response code", err)
      return
    }
    code = int(code)
    commands := strings.Split(line, " ")
    if code == 238 && len(commands) > 1 && ValidMessageID(commands[0]) {
      fname := d.store.GetFilename(commands[0])
      f, err := os.Open(fname)
      if f == nil {
        continue
      }
      err = self.SendMessage(commands[0], f, d)
      f.Close()
      if err != nil {
        log.Println("failed to send message", err)
        self.Quit()
        return
      }
    } else if code == 438 {
      // declined
      continue
    } else if code == 239 {
      // accepted
      continue
    } else if code == 439 {
      // invalid
      continue
    } else {
      log.Printf("invalid response from outbound feed: '%d %s'", code, line)
    }
  }
}

// just do it (tm)
func (self *NNTPConnection) SendMessage(msgid string, msg io.Reader, d *NNTPDaemon) error {
  var err error
  self.info.reading = true
  err = self.txtconn.PrintfLine("TAKETHIS %s", msgid)
  if err != nil {
    log.Println("error in outfeed", err)
    return  err
  }
  wr := self.txtconn.DotWriter()
  _, err = io.Copy(wr, msg)
  wr.Close()
  self.info.reading = false
  if err != nil {
    log.Printf("failed to send %s via feed: %s", msgid, err)
    return err
  }
  return nil
}

// handle inbound connection
func (self *NNTPConnection) HandleInbound(d *NNTPDaemon) {

  // intitiate handshake
  var err error
  self.info.mode = "STREAM"
  log.Println("Incoming nntp connection from", self.conn.RemoteAddr())
  // send welcome
  greet := "2nd generation overchan NNTP Daemon posting allowed"
  self.txtconn.PrintfLine("200 %s", greet)
  for {
    if err != nil {
      log.Println("failure in infeed", err)
      self.Quit()
      return
    }
    // read line and break if needed
    line, err := self.ReadLine()
    if len(line) == 0 || err != nil {
      break
    }
    var code int
    var msg string
    commands := strings.Split(line, " ")
    cmd := commands[0]
    // capabilities command
    if cmd == "CAPABILITIES" {
      self.sendCapabilities()
    } else if cmd == "MODE" { // mode switch
      if len(commands) == 2 {
        // get mode
        mode := strings.ToUpper(commands[1])
        // reader mode
        if mode == "READER" {
          self.info.mode = "reader"
          code = 201
          msg = "posting disallowed"
        } else if mode == "STREAM" {
          // mode stream
          self.info.mode = "stream"
          code = 203
          msg = "stream it"
        } else {
          // other modes not implemented
          code = 501
          msg = "mode not implemented"
        }
      } else {
        code = 500
        msg = "syntax error"
      }

      self.txtconn.PrintfLine("%d %s", code, msg)
    } else if self.info.mode == "stream" { // we are in stream mode
      if cmd == "TAKETHIS" {
        var newsgroup string
        if len(commands) == 2 {
          article := commands[1]
          if ValidMessageID(article) {
            code := 239
            file := d.store.CreateTempFile(article)
            ip_header := ""
            ip_banned := false
            headers_done := false
            read_more := true
            has_attachment := false
            is_signed := false
            message := "we are gud"
            for {
              var line string
              line, err = self.ReadLine()
              if err != nil {
                log.Println("error reading", article, err)
                break
              }

              if line == "" && ! headers_done {
                // headers done
                headers_done = true
                if len(ip_header) > 0 {
                  ip_banned, err = d.database.CheckEncIPBanned(ip_header)
                  if err == nil {
                    if ip_banned {
                      code = 438
                      read_more = false
                      message = "poster is banned"
                    }
                  } else {
                    log.Println("cannot check for banned encrypted ip", err)
                    // send it later
                    code = 431
                    message = "could be banned by us but we do not know"
                  }
                } else if self.allow_tor {
                  // do we want tor posts with attachments?
                  if has_attachment && ! self.allow_tor_attachments {
                    // this guy is banned
                    code = 438
                    read_more = false
                    message = "we do not take attachments from tor"
                  }
                } else {
                  // we don'e want it
                  code = 438
                  message = "we do not take anonymous posts"
                  read_more = false
                }
                if is_signed {
                  log.Println("we got a signed message")
                }
              }

              if line == "." {
                break
              } else if ! headers_done {
                lower_line := strings.ToLower(line)
                // newsgroup header
                if strings.HasPrefix(lower_line, "newsgroups: ") {
                  if len(newsgroup) == 0 {
                    newsgroup := line[12:]
                    if ! newsgroupValidFormat(newsgroup) {
                      // bad newsgroup
                      code = 438
                      message = "invalid newsgroup"
                    }
                  }
                } else if strings.HasPrefix(lower_line, "x-tor-poster: 1") {
                  if ! self.allow_tor {
                    // we don't want this post
                    code = 438
                    message = "we do not take anonymous posts"
                    read_more = false
                  }
                } else if strings.HasPrefix(lower_line, "x-encrypted-ip: ") {
                  ip_header = strings.Split(line, " ")[1]
                  ip_header = strings.Trim(line, " \t\r\n")
                } else if strings.HasPrefix(lower_line, "content-type: multipart") {
                  has_attachment = true
                } else if strings.HasPrefix(lower_line, "x-signature-ed25519-sha512: ") {
                  is_signed = true
                  has_attachment = true
                }
              }
              if read_more {
                file.Write([]byte(line))
                file.Write([]byte("\n"))
              }
            }
            file.Close()
            // tell them our result
            self.txtconn.PrintfLine("%d %s %s", code, article, message)
            // the send was good
            if code == 239 {
              log.Println(self.conn.RemoteAddr(), "got article", article)
              // inform daemon
              d.infeed_load <- article
            } else if code == 431 {
              // defer
              log.Println("defering receival, keeping article in tmp dir:", article)
            } else {
              // delete unaccepted article
              log.Println("did not accept", article)
              fname := d.store.GetTempFilename(article)
              DelFile(fname)
            }
          } else {
            self.txtconn.PrintfLine("438 %s invalid message-id", article)
          }
        }
      }
      // check command
      if cmd == "CHECK" {
        if len(commands) == 2 {
          // check syntax
          // send error if needed
          article := commands[1]
          if ! ValidMessageID(article) {
            self.txtconn.PrintfLine("501 bad message id")
            continue
          }
          // do we already have this article?
          if d.store.HasArticle(article) {
            // ya, we got it already
            // tell them to not send it
            self.txtconn.PrintfLine("438 %s we have this article", article)
          } else {
            // nope, we do not have it
            // tell them to send it
            self.txtconn.PrintfLine("238 %s we want this article please give it", article)
          }
        }
      }
    }
  }
  self.Close()
}

func (self *NNTPConnection) sendCapabilities() {
  wr := self.txtconn.DotWriter()
  io.WriteString(wr, "101 we can haz do things\r\n")
  io.WriteString(wr, "VERSION 2\r\n")
  io.WriteString(wr, "IMPLEMENTATION srndv2 better than SRNd\r\n")
  io.WriteString(wr, "STREAMING\r\n")
  io.WriteString(wr, "READER\r\n")
  wr.Close()
}

func (self *NNTPConnection) Quit() {
  if ! self.inbound {
    self.txtconn.PrintfLine("QUIT")
  }
  self.Close()
}

func (self *NNTPConnection) ReadLine() (string, error) {
  line, err := self.txtconn.ReadLine()
  if err != nil {
    log.Println("error reading line in feed", err)
    return "", err
  }
  if self.debug {
    log.Println(self.conn.RemoteAddr(), "recv line", line)
  }
  return line, nil
}

// close the connection
func (self *NNTPConnection) Close() {
  err := self.conn.Close()
  if err != nil {
    log.Println(self.conn.RemoteAddr(), err)
  }
  log.Println(self.conn.RemoteAddr(), "Closed Connection")
}
<|endoftext|>


<commit_before>package test

import (
	"time"
	"os"
	"io/ioutil"
	"⚛sdl"
	"spectrum"
	"spectrum/formats"
	"prettytest"
 )

 var (
	speccy *spectrum.Spectrum48k
	app    *spectrum.Application
	r renderer
 )

type renderer struct {
	appSurface *sdl.Surface
	speccySurface *spectrum.SDLScreen2x
	width, height int
}

func (r *renderer) render(speccyRects []sdl.Rect) {
	for _, rect := range speccyRects {
		r.appSurface.Blit(&rect, r.speccySurface.GetSurface(), &rect)
		// r.appSurface.UpdateRect(int32(rect.X+r.config.consoleX), int32(rect.Y+r.config.consoleY), uint32(rect.W), uint32(rect.H))
	}
	r.appSurface.Flip()
}

func initDisplay() {
	r.width = spectrum.TotalScreenWidth*2
	r.height = spectrum.TotalScreenHeight*2

	sdlScreen := spectrum.NewSDLScreen2x(app)
	speccy.CommandChannel <- spectrum.Cmd_AddDisplay{sdlScreen}
	r.speccySurface = sdlScreen

	r.appSurface = sdl.SetVideoMode(r.width, r.height, 32, 0)
}

type testSuite struct { prettytest.Suite }

// func (t *testSuite) beforeAll() {
// 	StartFullEmulation()
// }

// func (t *testSuite) afterAll() {
// 	app.RequestExit()
// 	<-app.HasTerminated
// }

func (t *testSuite) before() {
	StartFullEmulation()
}

func (t *testSuite) after() {
	app.RequestExit()
	<-app.HasTerminated
}

func StartFullEmulation() {
	var (
		err os.Error
		speccyRects []sdl.Rect
	)

	app = spectrum.NewApplication()

	speccy, err = spectrum.NewSpectrum48k(app, "testdata/48.rom")
	speccy.TapeDrive().NotifyLoadComplete = true

	if err != nil {
		panic(err)
	}

	if sdl.Init(sdl.INIT_VIDEO|sdl.INIT_AUDIO) != 0 {
		app.PrintfMsg("%s", sdl.GetError())
		app.RequestExit()
		<-app.HasTerminated
		sdl.Quit()
	}

	sdl.WM_SetCaption("GoSpeccy - ZX Spectrum Emulator - Test mode", "")

	initDisplay()

	audio, err := spectrum.NewSDLAudio(app)

	if err == nil {
		speccy.CommandChannel <- spectrum.Cmd_AddAudioReceiver{audio}
	} else {
		app.PrintfMsg("%s", err)
	}

	go speccy.EmulatorLoop()

	ticker := time.NewTicker(1e9/int64(60))

	go func() {
		for {
			select {
			case speccyRects = <-r.speccySurface.UpdatedRectsCh():
			case <-ticker.C: r.render(speccyRects)
			}
		}
	}()

	<-speccy.ROMLoaded()
}

func loadSnapshot(filename string) formats.Snapshot {
	_snapshot, err := formats.ReadProgram(filename)

	if err != nil {
		panic(err)
	}

	var snapshot formats.Snapshot
	var ok bool
	if snapshot, ok = _snapshot.(formats.Snapshot); !ok {
		panic("invalid type")
	}

	return snapshot
}

func assertScreenEqual(expected, actual formats.Snapshot) bool {
	// Doesn't compare screen attributes
	for address, actualValue := range actual.Memory()[:0x1800] {
		if expected.Memory()[address] != actualValue {
			return false
		}
	}
	return true
}

func assertStateEqual(expected, actual formats.Snapshot) bool {
	// Compare memory ignoring PC value
	for address, actualValue := range actual.Memory() {
		if expected.Memory()[address] != actualValue && address != 0xbf4a && address != 0xbf4b {
			return false
		}
	}

	// FIXME: Should compare also CPU and ULA states

	return true
}

func stateEqualTo(filename string) bool {
	return assertStateEqual(loadSnapshot(filename), speccy.Cpu.MakeSnapshot())
}

func screenEqualTo(filename string) bool {
	return assertScreenEqual(loadSnapshot(filename), speccy.Cpu.MakeSnapshot())
}

func saveSnapshot(filename string) {
	fullSnapshot := speccy.Cpu.MakeSnapshot()

	data, err := fullSnapshot.EncodeSNA()

	if err != nil {
		panic(err)
	}

	err = ioutil.WriteFile(filename, data, 0600)

	if err != nil {
		panic(err)
	}
}
<commit_msg>Remove unneeded beforeAll, afterAll from test.go<commit_after>package test

import (
	"time"
	"os"
	"io/ioutil"
	"⚛sdl"
	"spectrum"
	"spectrum/formats"
	"prettytest"
 )

 var (
	speccy *spectrum.Spectrum48k
	app    *spectrum.Application
	r renderer
 )

type renderer struct {
	appSurface *sdl.Surface
	speccySurface *spectrum.SDLScreen2x
	width, height int
}

func (r *renderer) render(speccyRects []sdl.Rect) {
	for _, rect := range speccyRects {
		r.appSurface.Blit(&rect, r.speccySurface.GetSurface(), &rect)
		// r.appSurface.UpdateRect(int32(rect.X+r.config.consoleX), int32(rect.Y+r.config.consoleY), uint32(rect.W), uint32(rect.H))
	}
	r.appSurface.Flip()
}

func initDisplay() {
	r.width = spectrum.TotalScreenWidth*2
	r.height = spectrum.TotalScreenHeight*2

	sdlScreen := spectrum.NewSDLScreen2x(app)
	speccy.CommandChannel <- spectrum.Cmd_AddDisplay{sdlScreen}
	r.speccySurface = sdlScreen

	r.appSurface = sdl.SetVideoMode(r.width, r.height, 32, 0)
}

type testSuite struct { prettytest.Suite }

func (t *testSuite) before() {
	StartFullEmulation()
}

func (t *testSuite) after() {
	app.RequestExit()
	<-app.HasTerminated
}

func StartFullEmulation() {
	var (
		err os.Error
		speccyRects []sdl.Rect
	)

	app = spectrum.NewApplication()

	speccy, err = spectrum.NewSpectrum48k(app, "testdata/48.rom")
	speccy.TapeDrive().NotifyLoadComplete = true

	if err != nil {
		panic(err)
	}

	if sdl.Init(sdl.INIT_VIDEO|sdl.INIT_AUDIO) != 0 {
		app.PrintfMsg("%s", sdl.GetError())
		app.RequestExit()
		<-app.HasTerminated
		sdl.Quit()
	}

	sdl.WM_SetCaption("GoSpeccy - ZX Spectrum Emulator - Test mode", "")

	initDisplay()

	audio, err := spectrum.NewSDLAudio(app)

	if err == nil {
		speccy.CommandChannel <- spectrum.Cmd_AddAudioReceiver{audio}
	} else {
		app.PrintfMsg("%s", err)
	}

	go speccy.EmulatorLoop()

	ticker := time.NewTicker(1e9/int64(60))

	go func() {
		for {
			select {
			case speccyRects = <-r.speccySurface.UpdatedRectsCh():
			case <-ticker.C: r.render(speccyRects)
			}
		}
	}()

	<-speccy.ROMLoaded()
}

func loadSnapshot(filename string) formats.Snapshot {
	_snapshot, err := formats.ReadProgram(filename)

	if err != nil {
		panic(err)
	}

	var snapshot formats.Snapshot
	var ok bool
	if snapshot, ok = _snapshot.(formats.Snapshot); !ok {
		panic("invalid type")
	}

	return snapshot
}

func assertScreenEqual(expected, actual formats.Snapshot) bool {
	// Doesn't compare screen attributes
	for address, actualValue := range actual.Memory()[:0x1800] {
		if expected.Memory()[address] != actualValue {
			return false
		}
	}
	return true
}

func assertStateEqual(expected, actual formats.Snapshot) bool {
	// Compare memory ignoring PC value
	for address, actualValue := range actual.Memory() {
		if expected.Memory()[address] != actualValue && address != 0xbf4a && address != 0xbf4b {
			return false
		}
	}

	// FIXME: Should compare also CPU and ULA states

	return true
}

func stateEqualTo(filename string) bool {
	return assertStateEqual(loadSnapshot(filename), speccy.Cpu.MakeSnapshot())
}

func screenEqualTo(filename string) bool {
	return assertScreenEqual(loadSnapshot(filename), speccy.Cpu.MakeSnapshot())
}

func saveSnapshot(filename string) {
	fullSnapshot := speccy.Cpu.MakeSnapshot()

	data, err := fullSnapshot.EncodeSNA()

	if err != nil {
		panic(err)
	}

	err = ioutil.WriteFile(filename, data, 0600)

	if err != nil {
		panic(err)
	}
}
<|endoftext|>


<commit_before>// MIT License
//
// Copyright (c) 2016-2017 Logimethods
//
// Permission is hereby granted, free of charge, to any person obtaining a copy
// of this software and associated documentation files (the "Software"), to deal
// in the Software without restriction, including without limitation the rights
// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
// copies of the Software, and to permit persons to whom the Software is
// furnished to do so, subject to the following conditions:
//
// The above copyright notice and this permission notice shall be included in all
// copies or substantial portions of the Software.
//
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
// SOFTWARE.

package main

import (
  "os"
  "log"
  "fmt"
  "time"
  "strings"

//  "bytes"
    "encoding/binary"
    "math"
    "strconv"
//    "unsafe"

  "github.com/nats-io/go-nats"
  "github.com/gocql/gocql"
)

const query = "INSERT INTO raw_data (" +
        "line, transformer, usagePoint, year, month, day, hour, minute, day_of_week, voltage, demand, " +
        "val3, val4, val5, val6, val7, val8, val9, val10, val11, val12) " +
        "VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)"

// http://docs.datastax.com/en/cql/3.3/cql/cql_using/useCounters.html
const increment = "UPDATE raw_data_count SET count = count + 1 WHERE slot = ?"

var Session *gocql.Session
var Cluster *gocql.ClusterConfig

func main() {
  log.Print("Welcome to the NATS to Cassandra Bridge")

  log.Print("The time is", time.Now())

  nats_uri := os.Getenv("NATS_URI")
  nc, _ := nats.Connect(nats_uri)

  log.Print("Subscribed to NATS: ", nats_uri)

  nats_subject := os.Getenv("NATS_SUBJECT")
  log.Print("NATS Subject: ", nats_subject)

  // CASSANDRA
  cassandra_url := os.Getenv("CASSANDRA_URL")
  log.Print("Cassandra URL: ", cassandra_url)

  // LOG LEVEL
  log_level := os.Getenv("LOG_LEVEL")
  log.Print("LOG LEVEL: ", log_level)

  // TASK_SLOT
  task_slot := os.Getenv("TASK_SLOT")
  log.Print("TASK SLOT: ", task_slot)

  // connect to the Cluster
  Cluster = gocql.NewCluster(cassandra_url)
  Cluster.Keyspace = "smartmeter"

  cluster_consistency := os.Getenv("CASSANDRA_INJECT_CONSISTENCY")
  //// https://github.com/gocql/gocql/commit/939af06b39e61c17b91b7c11720213246ece9b80#diff-38a0fa12b0c511105c3a411f0df3e318
  Cluster.Consistency = gocql.ParseConsistency(cluster_consistency)

  cluster_timeout := os.Getenv("CASSANDRA_TIMEOUT")
  if (cluster_timeout != "") {
    timeout, err := strconv.Atoi(cluster_timeout)
    if (err == nil) {
      Cluster.Timeout = time.Duration(timeout) * time.Millisecond
      log.Print("(Provided) CASSANDRA_TIMEOUT: ", Cluster.Timeout)
    } else {
      log.Panicf("Unable to parse CASSANDRA_TIMEOUT (%s) into int64", cluster_timeout)
    }
  } else {
    log.Print("(Default) CASSANDRA_TIMEOUT: ", Cluster.Timeout)
  }

  createSession()
  defer Session.Close()

  log.Print("Connected to Cassandra")

  // Simple Async Subscriber
  nc.QueueSubscribe(nats_subject, "cassandra_inject", func(m *nats.Msg) {
    go insertIntoCassandra(m, task_slot, log_level)
  })

  log.Print("Ready to store NATS messages into CASSANDRA")

  // To keep the App alive
  for {
    time.Sleep(30 * time.Second)
    // log.Print(time.Now())
  }

}

func createSession() {
  log.Print("Cluster.CreateSession()")
  var err error
  Session, err = Cluster.CreateSession()
  if (err != nil) {
    log.Fatalf("Could not connect to Cassandra Cluster %s", err)
  }
}

func insertIntoCassandra(m *nats.Msg, task_slot string, log_level string) {
  /*** Point ***/

  // https://www.dotnetperls.com/split-go
  subjects := strings.Split(m.Subject, ".")
  // smartmeter.voltage.data.3.3.2   | (2016-11-16T20:05:04,116.366646)
  len := len(subjects)

  // http://stackoverflow.com/questions/30299649/golang-converting-string-to-specific-type-of-int-int8-int16-int32-int64
  line, _ := strconv.ParseInt(subjects[len -3], 10, 8) // tinyint,    // 8-bit signed int
  transformer, _ := strconv.ParseInt(subjects[len -2], 10, 32)  //  int, // 32-bit signed int
  usagePoint, _ := strconv.ParseInt(subjects[len -1], 10, 32)  //  int,

  /*** Date ***/

  longBytes := m.Data[:8]
  // http://stackoverflow.com/questions/22491876/convert-byte-array-uint8-to-float64-in-golang
  epoch := int64(binary.BigEndian.Uint64(longBytes))
  date := time.Unix(epoch, 0)

  // https://golang.org/pkg/time/#Time
  year, month, day := date.Date()
  hour, minute, _ := date.Clock()
  day_of_week := date.Weekday()

  /*** Voltage ***/

  voltageFloatBytes := m.Data[8:12]
  // http://stackoverflow.com/questions/22491876/convert-byte-array-uint8-to-float64-in-golang
  voltage := math.Float32frombits(binary.BigEndian.Uint32(voltageFloatBytes))

  /*** Demand ***/

  demandFloatBytes := m.Data[12:16]
  // http://stackoverflow.com/questions/22491876/convert-byte-array-uint8-to-float64-in-golang
  demand := math.Float32frombits(binary.BigEndian.Uint32(demandFloatBytes))

  /*** Remaining Values ***/

  var values [10]float32
  for i := 0; i < 10; i++ {
    values[i] = math.Float32frombits(binary.BigEndian.Uint32(m.Data[(16+4*i):(20+4*i)]))
  }

  // fmt.Print(".")

  if (log_level == "TRACE") {
    s := fmt.Sprintf("- v: %d", voltage)
    log.Print(s)
  }

  if (Session == nil || Session.Closed()) {
    createSession()
  }

  /** insert the Data into Cassandra **/

  if err := Session.Query(query,
      int8(line), int32(transformer), int32(usagePoint), int16(year), int8(month), int8(day),
      int8(hour), int8(minute), int8(day_of_week), voltage, demand,
      values[0], values[1], values[2], values[3], values[4],
      values[5], values[6], values[7], values[8], values[9]).Exec(); err != nil {
      log.Panic(err)
  } else {
    if err := Session.Query(increment, task_slot).Exec(); err != nil {
        log.Panic(err)
    }
  }
}
<commit_msg>"gopkg.in/nats-io/go-nats.v1"<commit_after>// MIT License
//
// Copyright (c) 2016-2017 Logimethods
//
// Permission is hereby granted, free of charge, to any person obtaining a copy
// of this software and associated documentation files (the "Software"), to deal
// in the Software without restriction, including without limitation the rights
// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
// copies of the Software, and to permit persons to whom the Software is
// furnished to do so, subject to the following conditions:
//
// The above copyright notice and this permission notice shall be included in all
// copies or substantial portions of the Software.
//
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
// SOFTWARE.

package main

import (
  "os"
  "log"
  "fmt"
  "time"
  "strings"

//  "bytes"
  "encoding/binary"
  "math"
  "strconv"
//    "unsafe"

  // http://labix.org/gopkg.in
  "gopkg.in/nats-io/go-nats.v1"
  "github.com/gocql/gocql"
)

const query = "INSERT INTO raw_data (" +
        "line, transformer, usagePoint, year, month, day, hour, minute, day_of_week, voltage, demand, " +
        "val3, val4, val5, val6, val7, val8, val9, val10, val11, val12) " +
        "VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)"

// http://docs.datastax.com/en/cql/3.3/cql/cql_using/useCounters.html
const increment = "UPDATE raw_data_count SET count = count + 1 WHERE slot = ?"

var Session *gocql.Session
var Cluster *gocql.ClusterConfig

func main() {
  log.Print("Welcome to the NATS to Cassandra Bridge")

  log.Print("The time is", time.Now())

  nats_uri := os.Getenv("NATS_URI")
  nc, _ := nats.Connect(nats_uri)

  log.Print("Subscribed to NATS: ", nats_uri)

  nats_subject := os.Getenv("NATS_SUBJECT")
  log.Print("NATS Subject: ", nats_subject)

  // CASSANDRA
  cassandra_url := os.Getenv("CASSANDRA_URL")
  log.Print("Cassandra URL: ", cassandra_url)

  // LOG LEVEL
  log_level := os.Getenv("LOG_LEVEL")
  log.Print("LOG LEVEL: ", log_level)

  // TASK_SLOT
  task_slot := os.Getenv("TASK_SLOT")
  log.Print("TASK SLOT: ", task_slot)

  // connect to the Cluster
  Cluster = gocql.NewCluster(cassandra_url)
  Cluster.Keyspace = "smartmeter"

  cluster_consistency := os.Getenv("CASSANDRA_INJECT_CONSISTENCY")
  //// https://github.com/gocql/gocql/commit/939af06b39e61c17b91b7c11720213246ece9b80#diff-38a0fa12b0c511105c3a411f0df3e318
  Cluster.Consistency = gocql.ParseConsistency(cluster_consistency)

  cluster_timeout := os.Getenv("CASSANDRA_TIMEOUT")
  if (cluster_timeout != "") {
    timeout, err := strconv.Atoi(cluster_timeout)
    if (err == nil) {
      Cluster.Timeout = time.Duration(timeout) * time.Millisecond
      log.Print("(Provided) CASSANDRA_TIMEOUT: ", Cluster.Timeout)
    } else {
      log.Panicf("Unable to parse CASSANDRA_TIMEOUT (%s) into int64", cluster_timeout)
    }
  } else {
    log.Print("(Default) CASSANDRA_TIMEOUT: ", Cluster.Timeout)
  }

  createSession()
  defer Session.Close()

  log.Print("Connected to Cassandra")

  // Simple Async Subscriber
  nc.QueueSubscribe(nats_subject, "cassandra_inject", func(m *nats.Msg) {
    go insertIntoCassandra(m, task_slot, log_level)
  })

  log.Print("Ready to store NATS messages into CASSANDRA")

  // To keep the App alive
  for {
    time.Sleep(30 * time.Second)
    // log.Print(time.Now())
  }

}

func createSession() {
  log.Print("Cluster.CreateSession()")
  var err error
  Session, err = Cluster.CreateSession()
  if (err != nil) {
    log.Fatalf("Could not connect to Cassandra Cluster %s", err)
  }
}

func insertIntoCassandra(m *nats.Msg, task_slot string, log_level string) {
  /*** Point ***/

  // https://www.dotnetperls.com/split-go
  subjects := strings.Split(m.Subject, ".")
  // smartmeter.voltage.data.3.3.2   | (2016-11-16T20:05:04,116.366646)
  len := len(subjects)

  // http://stackoverflow.com/questions/30299649/golang-converting-string-to-specific-type-of-int-int8-int16-int32-int64
  line, _ := strconv.ParseInt(subjects[len -3], 10, 8) // tinyint,    // 8-bit signed int
  transformer, _ := strconv.ParseInt(subjects[len -2], 10, 32)  //  int, // 32-bit signed int
  usagePoint, _ := strconv.ParseInt(subjects[len -1], 10, 32)  //  int,

  /*** Date ***/

  longBytes := m.Data[:8]
  // http://stackoverflow.com/questions/22491876/convert-byte-array-uint8-to-float64-in-golang
  epoch := int64(binary.BigEndian.Uint64(longBytes))
  date := time.Unix(epoch, 0)

  // https://golang.org/pkg/time/#Time
  year, month, day := date.Date()
  hour, minute, _ := date.Clock()
  day_of_week := date.Weekday()

  /*** Voltage ***/

  voltageFloatBytes := m.Data[8:12]
  // http://stackoverflow.com/questions/22491876/convert-byte-array-uint8-to-float64-in-golang
  voltage := math.Float32frombits(binary.BigEndian.Uint32(voltageFloatBytes))

  /*** Demand ***/

  demandFloatBytes := m.Data[12:16]
  // http://stackoverflow.com/questions/22491876/convert-byte-array-uint8-to-float64-in-golang
  demand := math.Float32frombits(binary.BigEndian.Uint32(demandFloatBytes))

  /*** Remaining Values ***/

  var values [10]float32
  for i := 0; i < 10; i++ {
    values[i] = math.Float32frombits(binary.BigEndian.Uint32(m.Data[(16+4*i):(20+4*i)]))
  }

  // fmt.Print(".")

  if (log_level == "TRACE") {
    s := fmt.Sprintf("- v: %d", voltage)
    log.Print(s)
  }

  if (Session == nil || Session.Closed()) {
    createSession()
  }

  /** insert the Data into Cassandra **/

  if err := Session.Query(query,
      int8(line), int32(transformer), int32(usagePoint), int16(year), int8(month), int8(day),
      int8(hour), int8(minute), int8(day_of_week), voltage, demand,
      values[0], values[1], values[2], values[3], values[4],
      values[5], values[6], values[7], values[8], values[9]).Exec(); err != nil {
      log.Panic(err)
  } else {
    if err := Session.Query(increment, task_slot).Exec(); err != nil {
        log.Panic(err)
    }
  }
}
<|endoftext|>


<commit_before>/*
   Nging is a toolbox for webmasters
   Copyright (C) 2018-present  Wenhui Shen <swh@admpub.com>

   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 <https://www.gnu.org/licenses/>.
*/

package config

import (
	"errors"
	"io"
	"os"
	"os/exec"
	"path/filepath"
	"strings"
	"sync"
	"time"

	"github.com/admpub/log"
	"github.com/admpub/nging/v4/application/cmd/event"
	"github.com/admpub/nging/v4/application/library/config/cmder"
	"github.com/admpub/nging/v4/application/library/config/startup"
	"github.com/spf13/pflag"
	"github.com/webx-top/com"
	"github.com/webx-top/echo"
	"github.com/webx-top/echo/param"
)

func NewCLIConfig() *CLIConfig {
	cli := &CLIConfig{
		cmds:     map[string]*exec.Cmd{},
		pid:      os.Getpid(),
		envFiles: findEnvFile(),
	}
	cli.InitEnviron()
	//cli.WatchEnvConfig()
	return cli
}

var DefaultStartup = `webserver,task,daemon,ftpserver,frpserver,frpclient`
var DefaultPort = 9999

type CLIConfig struct {
	BackendDomain  string //后台绑定域名
	FrontendDomain string //前台绑定域名
	Address        string //监听IP地址
	Port           int    //监听端口
	Conf           string
	Confx          string
	Type           string //启动类型: webserver/ftpserver/manager
	Startup        string //manager启动时同时启动的服务，可选的有webserver/ftpserver,如有多个需用半角逗号“,”隔开
	cmds           map[string]*exec.Cmd
	pid            int
	envVars        map[string]string // 从文件“.env”中读取到的自定义环境变量（对于系统环境变量中已经存在的变量，会自动忽略）
	envFiles       []string
	envMonitor     *com.MonitorEvent
	envLock        sync.RWMutex
}

func (c *CLIConfig) InitFlag(flagSet *pflag.FlagSet) {
	flagSet.StringVarP(&c.Address, `address`, `a`, `0.0.0.0`, `address`)
	flagSet.IntVarP(&c.Port, `port`, `p`, DefaultPort, `port`)
	flagSet.StringVarP(&c.Conf, `config`, `c`, filepath.Join(echo.Wd(), `config/config.yaml`), `config`)
	flagSet.StringVarP(&c.Confx, `subconfig`, `u`, filepath.Join(echo.Wd(), `config/config.frpserver.yaml`), `submodule config`)
	flagSet.StringVarP(&c.Type, `type`, `t`, `manager`, `operation type`)
	flagSet.StringVarP(&c.Startup, `startup`, `s`, DefaultStartup, `startup`)
	flagSet.StringVarP(&c.FrontendDomain, `frontend.domain`, `f`, ``, `frontend domain`)
	flagSet.StringVarP(&c.BackendDomain, `backend.domain`, `b`, ``, `backend domain`)
}

func (c *CLIConfig) OnlyRunServer() bool {
	cm := cmder.Get(c.Type)
	if cm != nil {
		startup.FireBefore(c.Type)
		err := cm.Init()
		if err != nil {
			com.ExitOnFailure(err.Error())
		}
		startup.FireAfter(c.Type)
		return true
	}

	// manager mode
	if c.Type == `official` || !event.SupportManager {
		c.Startup = `none`
	}
	return false
}

func (c *CLIConfig) ParseConfig() {
	err := ParseConfig()
	if err != nil {
		if os.IsNotExist(err) {
			panic(err)
		}
		if IsInstalled() {
			MustOK(err)
		} else {
			log.Error(err)
		}
	}
}

//RunStartup manager启动时同时启动的服务
func (c *CLIConfig) RunStartup() {
	c.ParseConfig()
	c.Startup = strings.TrimSpace(c.Startup)
	if len(c.Startup) < 1 || !IsInstalled() {
		return
	}
	for _, serverType := range param.StringSlice(strings.Split(c.Startup, `,`)).Unique() {
		serverType = strings.TrimSpace(serverType)
		cm := cmder.Get(serverType)
		if cm != nil {
			startup.FireBefore(serverType)
			if err := cm.Restart(); err != nil {
				log.Error(err)
			}
			startup.FireAfter(serverType)
			continue
		}
	}
}

func (c *CLIConfig) CmdGroupStop(groupName string) error {
	if c.cmds == nil {
		return nil
	}
	groupName += `.`
	var err error
	for key, cmd := range c.cmds {
		if !strings.HasPrefix(key, groupName) {
			continue
		}
		err := c.Kill(cmd)
		if err != nil {
			log.Error(err)
		}
	}
	return err
}

func (c *CLIConfig) CmdHasGroup(groupName string) bool {
	if c.cmds == nil {
		return false
	}
	groupName += `.`
	for key := range c.cmds {
		if !strings.HasPrefix(key, groupName) {
			continue
		}
		return true
	}
	return false
}

func (c *CLIConfig) CmdGet(typeName string) *exec.Cmd {
	if c.cmds == nil {
		return nil
	}
	cmd, _ := c.cmds[typeName]
	return cmd
}

func (c *CLIConfig) CmdSet(name string, cmd *exec.Cmd) {
	c.cmds[name] = cmd
}

func (c *CLIConfig) CmdStop(typeName string) error {
	cmd := c.CmdGet(typeName)
	if cmd == nil {
		return nil
	}
	return c.Kill(cmd)
}

func (c *CLIConfig) CmdSendSignal(typeName string, sig os.Signal) error {
	cmd := c.CmdGet(typeName)
	return c.SendSignal(cmd, sig)
}

func (c *CLIConfig) SendSignal(cmd *exec.Cmd, sig os.Signal) error {
	if cmd == nil {
		return nil
	}
	if cmd.Process == nil {
		return nil
	}
	err := cmd.Process.Signal(sig)
	if err != nil && (cmd.ProcessState == nil || cmd.ProcessState.Exited()) {
		err = nil
	}
	return err
}

func (c *CLIConfig) Kill(cmd *exec.Cmd) error {
	return com.CloseProcessFromCmd(cmd)
}

var ErrCmdNotRunning = errors.New(`command is not running`)

func (c *CLIConfig) SetLogWriter(cmdType string, writer ...io.Writer) error {
	if c.cmds == nil {
		return ErrCmdNotRunning
	}
	cmd, ok := c.cmds[cmdType]
	if !ok || cmd == nil {
		return ErrCmdNotRunning
	}
	var wOut, wErr io.Writer
	length := len(writer)
	if length > 0 {
		wOut = writer[0]
		if length > 1 {
			wErr = writer[1]
		} else {
			wErr = wOut
		}
	} else {
		wOut = os.Stdout
		wErr = os.Stderr
	}
	cmd.Stdout = wOut
	cmd.Stderr = wErr
	return nil
}

func (c *CLIConfig) IsRunning(ct string) bool {
	if c.cmds == nil {
		return false
	}
	cmd, ok := c.cmds[ct]
	if !ok {
		return false
	}
	return CmdIsRunning(cmd)
}

type ReloadByConfig interface {
	ReloadByConfig(cfg *Config, args ...string) error
}

func (c *CLIConfig) Reload(cfg *Config, cts ...string) error {
	for _, ct := range cts {
		if !c.IsRunning(ct) {
			continue
		}
		typeAndID := strings.SplitN(ct, ".", 2)
		serverType := typeAndID[0]
		cm := cmder.Get(serverType)
		if cm != nil {
			if rd, ok := cm.(ReloadByConfig); ok {
				if len(typeAndID) == 2 {
					rd.ReloadByConfig(cfg, typeAndID[1])
					continue
				}
				rd.ReloadByConfig(cfg)
				continue
			}
			if len(typeAndID) == 1 {
				cm.Reload()
				continue
			}
			if rd, ok := cm.(cmder.RestartBy); ok {
				if len(typeAndID) == 2 {
					rd.RestartBy(typeAndID[1])
					continue
				}
			}
			cm.Reload()
			continue
		}
	}
	return nil
}

func (c *CLIConfig) GenerateIDFromConfigFileName(configFile string, musts ...bool) string {
	baseName := filepath.Base(configFile)
	index := strings.LastIndex(baseName, `.`)
	var id string
	if index > 0 {
		id = baseName[0:index]
	}
	if len(musts) == 0 || !musts[0] {
		if len(id) == 0 {
			id = time.Now().Format(`020060102150405`)
		}
	}
	return id
}

func (c *CLIConfig) Close() error {
	for _, cmd := range c.cmds {
		err := c.Kill(cmd)
		if err != nil {
			log.Error(err)
		}
	}
	return nil
}

func (c *CLIConfig) Pid() int {
	return c.pid
}
<commit_msg>update<commit_after>/*
   Nging is a toolbox for webmasters
   Copyright (C) 2018-present  Wenhui Shen <swh@admpub.com>

   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 <https://www.gnu.org/licenses/>.
*/

package config

import (
	"errors"
	"io"
	"os"
	"os/exec"
	"path/filepath"
	"strings"
	"sync"
	"time"

	"github.com/admpub/log"
	"github.com/admpub/nging/v4/application/cmd/event"
	"github.com/admpub/nging/v4/application/library/config/cmder"
	"github.com/admpub/nging/v4/application/library/config/startup"
	"github.com/spf13/pflag"
	"github.com/webx-top/com"
	"github.com/webx-top/echo"
	"github.com/webx-top/echo/param"
)

func NewCLIConfig() *CLIConfig {
	cli := &CLIConfig{
		cmds:     map[string]*exec.Cmd{},
		pid:      os.Getpid(),
		envFiles: findEnvFile(),
	}
	cli.InitEnviron()
	//cli.WatchEnvConfig()
	return cli
}

var DefaultStartup = `task`
var DefaultPort = 9999

type CLIConfig struct {
	BackendDomain  string //后台绑定域名
	FrontendDomain string //前台绑定域名
	Address        string //监听IP地址
	Port           int    //监听端口
	Conf           string
	Confx          string
	Type           string //启动类型: webserver/ftpserver/manager
	Startup        string //manager启动时同时启动的服务，可选的有webserver/ftpserver,如有多个需用半角逗号“,”隔开
	cmds           map[string]*exec.Cmd
	pid            int
	envVars        map[string]string // 从文件“.env”中读取到的自定义环境变量（对于系统环境变量中已经存在的变量，会自动忽略）
	envFiles       []string
	envMonitor     *com.MonitorEvent
	envLock        sync.RWMutex
}

func (c *CLIConfig) InitFlag(flagSet *pflag.FlagSet) {
	flagSet.StringVarP(&c.Address, `address`, `a`, `0.0.0.0`, `address`)
	flagSet.IntVarP(&c.Port, `port`, `p`, DefaultPort, `port`)
	flagSet.StringVarP(&c.Conf, `config`, `c`, filepath.Join(echo.Wd(), `config/config.yaml`), `config`)
	flagSet.StringVarP(&c.Confx, `subconfig`, `u`, filepath.Join(echo.Wd(), `config/config.frpserver.yaml`), `submodule config`)
	flagSet.StringVarP(&c.Type, `type`, `t`, `manager`, `operation type`)
	flagSet.StringVarP(&c.Startup, `startup`, `s`, DefaultStartup, `startup`)
	flagSet.StringVarP(&c.FrontendDomain, `frontend.domain`, `f`, ``, `frontend domain`)
	flagSet.StringVarP(&c.BackendDomain, `backend.domain`, `b`, ``, `backend domain`)
}

func (c *CLIConfig) OnlyRunServer() bool {
	cm := cmder.Get(c.Type)
	if cm != nil {
		startup.FireBefore(c.Type)
		err := cm.Init()
		if err != nil {
			com.ExitOnFailure(err.Error())
		}
		startup.FireAfter(c.Type)
		return true
	}

	// manager mode
	if c.Type == `official` || !event.SupportManager {
		c.Startup = `none`
	}
	return false
}

func (c *CLIConfig) ParseConfig() {
	err := ParseConfig()
	if err != nil {
		if os.IsNotExist(err) {
			panic(err)
		}
		if IsInstalled() {
			MustOK(err)
		} else {
			log.Error(err)
		}
	}
}

//RunStartup manager启动时同时启动的服务
func (c *CLIConfig) RunStartup() {
	c.ParseConfig()
	c.Startup = strings.TrimSpace(c.Startup)
	if len(c.Startup) < 1 || !IsInstalled() {
		return
	}
	for _, serverType := range param.StringSlice(strings.Split(c.Startup, `,`)).Unique() {
		serverType = strings.TrimSpace(serverType)
		cm := cmder.Get(serverType)
		if cm != nil {
			startup.FireBefore(serverType)
			if err := cm.Restart(); err != nil {
				log.Error(err)
			}
			startup.FireAfter(serverType)
			continue
		}
	}
}

func (c *CLIConfig) CmdGroupStop(groupName string) error {
	if c.cmds == nil {
		return nil
	}
	groupName += `.`
	var err error
	for key, cmd := range c.cmds {
		if !strings.HasPrefix(key, groupName) {
			continue
		}
		err := c.Kill(cmd)
		if err != nil {
			log.Error(err)
		}
	}
	return err
}

func (c *CLIConfig) CmdHasGroup(groupName string) bool {
	if c.cmds == nil {
		return false
	}
	groupName += `.`
	for key := range c.cmds {
		if !strings.HasPrefix(key, groupName) {
			continue
		}
		return true
	}
	return false
}

func (c *CLIConfig) CmdGet(typeName string) *exec.Cmd {
	if c.cmds == nil {
		return nil
	}
	cmd, _ := c.cmds[typeName]
	return cmd
}

func (c *CLIConfig) CmdSet(name string, cmd *exec.Cmd) {
	c.cmds[name] = cmd
}

func (c *CLIConfig) CmdStop(typeName string) error {
	cmd := c.CmdGet(typeName)
	if cmd == nil {
		return nil
	}
	return c.Kill(cmd)
}

func (c *CLIConfig) CmdSendSignal(typeName string, sig os.Signal) error {
	cmd := c.CmdGet(typeName)
	return c.SendSignal(cmd, sig)
}

func (c *CLIConfig) SendSignal(cmd *exec.Cmd, sig os.Signal) error {
	if cmd == nil {
		return nil
	}
	if cmd.Process == nil {
		return nil
	}
	err := cmd.Process.Signal(sig)
	if err != nil && (cmd.ProcessState == nil || cmd.ProcessState.Exited()) {
		err = nil
	}
	return err
}

func (c *CLIConfig) Kill(cmd *exec.Cmd) error {
	return com.CloseProcessFromCmd(cmd)
}

var ErrCmdNotRunning = errors.New(`command is not running`)

func (c *CLIConfig) SetLogWriter(cmdType string, writer ...io.Writer) error {
	if c.cmds == nil {
		return ErrCmdNotRunning
	}
	cmd, ok := c.cmds[cmdType]
	if !ok || cmd == nil {
		return ErrCmdNotRunning
	}
	var wOut, wErr io.Writer
	length := len(writer)
	if length > 0 {
		wOut = writer[0]
		if length > 1 {
			wErr = writer[1]
		} else {
			wErr = wOut
		}
	} else {
		wOut = os.Stdout
		wErr = os.Stderr
	}
	cmd.Stdout = wOut
	cmd.Stderr = wErr
	return nil
}

func (c *CLIConfig) IsRunning(ct string) bool {
	if c.cmds == nil {
		return false
	}
	cmd, ok := c.cmds[ct]
	if !ok {
		return false
	}
	return CmdIsRunning(cmd)
}

type ReloadByConfig interface {
	ReloadByConfig(cfg *Config, args ...string) error
}

func (c *CLIConfig) Reload(cfg *Config, cts ...string) error {
	for _, ct := range cts {
		if !c.IsRunning(ct) {
			continue
		}
		typeAndID := strings.SplitN(ct, ".", 2)
		serverType := typeAndID[0]
		cm := cmder.Get(serverType)
		if cm != nil {
			if rd, ok := cm.(ReloadByConfig); ok {
				if len(typeAndID) == 2 {
					rd.ReloadByConfig(cfg, typeAndID[1])
					continue
				}
				rd.ReloadByConfig(cfg)
				continue
			}
			if len(typeAndID) == 1 {
				cm.Reload()
				continue
			}
			if rd, ok := cm.(cmder.RestartBy); ok {
				if len(typeAndID) == 2 {
					rd.RestartBy(typeAndID[1])
					continue
				}
			}
			cm.Reload()
			continue
		}
	}
	return nil
}

func (c *CLIConfig) GenerateIDFromConfigFileName(configFile string, musts ...bool) string {
	baseName := filepath.Base(configFile)
	index := strings.LastIndex(baseName, `.`)
	var id string
	if index > 0 {
		id = baseName[0:index]
	}
	if len(musts) == 0 || !musts[0] {
		if len(id) == 0 {
			id = time.Now().Format(`020060102150405`)
		}
	}
	return id
}

func (c *CLIConfig) Close() error {
	for _, cmd := range c.cmds {
		err := c.Kill(cmd)
		if err != nil {
			log.Error(err)
		}
	}
	return nil
}

func (c *CLIConfig) Pid() int {
	return c.pid
}
<|endoftext|>


<commit_before>//go:generate ../scripts/generate-bindings bindings.js

package stitch

import (
	"crypto/sha1"
	"encoding/json"
	"fmt"
	"path/filepath"

	"github.com/robertkrimen/otto"

	// Automatically import the Javascript underscore utility-belt library into
	// the Stitch VM.
	_ "github.com/robertkrimen/otto/underscore"

	"github.com/quilt/quilt/util"
)

// A Stitch is an abstract representation of the policy language.
type Stitch struct {
	Containers  []Container  `json:",omitempty"`
	Labels      []Label      `json:",omitempty"`
	Connections []Connection `json:",omitempty"`
	Placements  []Placement  `json:",omitempty"`
	Machines    []Machine    `json:",omitempty"`

	AdminACL  []string `json:",omitempty"`
	MaxPrice  float64  `json:",omitempty"`
	Namespace string   `json:",omitempty"`

	Invariants []invariant `json:",omitempty"`
}

// A Placement constraint guides where containers may be scheduled, either relative to
// the labels of other containers, or the machine the container will run on.
type Placement struct {
	TargetLabel string `json:",omitempty"`

	Exclusive bool `json:",omitempty"`

	// Label Constraint
	OtherLabel string `json:",omitempty"`

	// Machine Constraints
	Provider   string `json:",omitempty"`
	Size       string `json:",omitempty"`
	Region     string `json:",omitempty"`
	FloatingIP string `json:",omitempty"`
}

// An Image represents a Docker image that can be run. If the Dockerfile is non-empty,
// the image should be built and hosted by Quilt.
type Image struct {
	Name       string `json:",omitempty"`
	Dockerfile string `json:",omitempty"`
}

// A Container may be instantiated in the stitch and queried by users.
type Container struct {
	ID                string            `json:",omitempty"`
	Image             Image             `json:",omitempty"`
	Command           []string          `json:",omitempty"`
	Env               map[string]string `json:",omitempty"`
	FilepathToContent map[string]string `json:",omitempty"`
	Hostname          string            `json:",omitempty"`
}

// A Label represents a logical group of containers.
type Label struct {
	Name        string   `json:",omitempty"`
	IDs         []string `json:",omitempty"`
	Annotations []string `json:",omitempty"`
}

// A Connection allows containers implementing the From label to speak to containers
// implementing the To label in ports in the range [MinPort, MaxPort]
type Connection struct {
	From    string `json:",omitempty"`
	To      string `json:",omitempty"`
	MinPort int    `json:",omitempty"`
	MaxPort int    `json:",omitempty"`
}

// A ConnectionSlice allows for slices of Collections to be used in joins
type ConnectionSlice []Connection

// A Machine specifies the type of VM that should be booted.
type Machine struct {
	ID         string   `json:",omitempty"`
	Provider   string   `json:",omitempty"`
	Role       string   `json:",omitempty"`
	Size       string   `json:",omitempty"`
	CPU        Range    `json:",omitempty"`
	RAM        Range    `json:",omitempty"`
	DiskSize   int      `json:",omitempty"`
	Region     string   `json:",omitempty"`
	SSHKeys    []string `json:",omitempty"`
	FloatingIP string   `json:",omitempty"`
}

// A Range defines a range of acceptable values for a Machine attribute
type Range struct {
	Min float64 `json:",omitempty"`
	Max float64 `json:",omitempty"`
}

// PublicInternetLabel is a magic label that allows connections to or from the public
// network.
const PublicInternetLabel = "public"

// Accepts returns true if `x` is within the range specified by `stitchr` (include),
// or if no max is specified and `x` is larger than `stitchr.min`.
func (stitchr Range) Accepts(x float64) bool {
	return stitchr.Min <= x && (stitchr.Max == 0 || x <= stitchr.Max)
}

func run(vm *otto.Otto, filename string, code string) (otto.Value, error) {
	// Compile before running so that stacktraces have filenames.
	script, err := vm.Compile(filename, code)
	if err != nil {
		return otto.Value{}, err
	}

	return vm.Run(script)
}

func newVM(getter ImportGetter) (*otto.Otto, error) {
	vm := otto.New()
	if err := vm.Set("githubKeys", toOttoFunc(githubKeysImpl)); err != nil {
		return vm, err
	}
	if err := vm.Set("require", toOttoFunc(getter.requireImpl)); err != nil {
		return vm, err
	}
	if err := vm.Set("hash", toOttoFunc(hashImpl)); err != nil {
		return vm, err
	}
	if err := vm.Set("read", toOttoFunc(readImpl)); err != nil {
		return vm, err
	}

	_, err := run(vm, "<javascript_bindings>", javascriptBindings)
	return vm, err
}

// `runSpec` evaluates `spec` within a module closure.
func runSpec(vm *otto.Otto, filename string, spec string) (otto.Value, error) {
	// The function declaration must be prepended to the first line of the
	// import or else stacktraces will show an offset line number.
	exec := "(function() {" +
		"var module={exports: {}};" +
		"(function(module, exports) {" +
		spec +
		"})(module, module.exports);" +
		"return module.exports" +
		"})()"
	return run(vm, filename, exec)
}

// New parses and executes a stitch (in text form), and returns an abstract Dsl handle.
func New(filename string, specStr string, getter ImportGetter) (Stitch, error) {
	vm, err := newVM(getter)
	if err != nil {
		return Stitch{}, err
	}

	if _, err := runSpec(vm, filename, specStr); err != nil {
		return Stitch{}, err
	}

	spec, err := parseContext(vm)
	if err != nil {
		return Stitch{}, err
	}
	spec.createPortRules()

	if len(spec.Invariants) == 0 {
		return spec, nil
	}

	graph, err := InitializeGraph(spec)
	if err != nil {
		return Stitch{}, err
	}

	if err := checkInvariants(graph, spec.Invariants); err != nil {
		return Stitch{}, err
	}

	return spec, nil
}

// FromJavascript gets a Stitch handle from a string containing Javascript code.
func FromJavascript(specStr string, getter ImportGetter) (Stitch, error) {
	return New("<raw_string>", specStr, getter)
}

// FromFile gets a Stitch handle from a file on disk.
func FromFile(filename string, getter ImportGetter) (Stitch, error) {
	specStr, err := util.ReadFile(filename)
	if err != nil {
		return Stitch{}, err
	}
	return New(filename, specStr, getter)
}

// FromJSON gets a Stitch handle from the deployment representation.
func FromJSON(jsonStr string) (stc Stitch, err error) {
	err = json.Unmarshal([]byte(jsonStr), &stc)
	return stc, err
}

func parseContext(vm *otto.Otto) (stc Stitch, err error) {
	vmCtx, err := vm.Run("deployment.toQuiltRepresentation()")
	if err != nil {
		return stc, err
	}

	// Export() always returns `nil` as the error (it's only present for
	// backwards compatibility), so we can safely ignore it.
	exp, _ := vmCtx.Export()
	ctxStr, err := json.Marshal(exp)
	if err != nil {
		return stc, err
	}
	err = json.Unmarshal(ctxStr, &stc)
	return stc, err
}

// createPortRules creates exclusive placement rules such that no two containers
// listening on the same public port get placed on the same machine.
func (stitch *Stitch) createPortRules() {
	ports := make(map[int][]string)
	for _, c := range stitch.Connections {
		if c.From != PublicInternetLabel && c.To != PublicInternetLabel {
			continue
		}

		target := c.From
		if c.From == PublicInternetLabel {
			target = c.To
		}

		min := c.MinPort
		ports[min] = append(ports[min], target)
	}

	for _, labels := range ports {
		for _, tgt := range labels {
			for _, other := range labels {
				stitch.Placements = append(stitch.Placements,
					Placement{
						Exclusive:   true,
						TargetLabel: tgt,
						OtherLabel:  other,
					})
			}
		}
	}
}

// String returns the Stitch in its deployment representation.
func (stitch Stitch) String() string {
	jsonBytes, err := json.Marshal(stitch)
	if err != nil {
		panic(err)
	}
	return string(jsonBytes)
}

func hashImpl(call otto.FunctionCall) (otto.Value, error) {
	if len(call.ArgumentList) < 1 {
		panic(call.Otto.MakeRangeError(
			"hash requires an argument"))
	}

	toHash, err := call.Argument(0).ToString()
	if err != nil {
		return otto.Value{}, err
	}

	return call.Otto.ToValue(fmt.Sprintf("%x", sha1.Sum([]byte(toHash))))
}

func readImpl(call otto.FunctionCall) (otto.Value, error) {
	if len(call.ArgumentList) < 1 {
		panic(call.Otto.MakeRangeError(
			"read requires an argument"))
	}

	path, err := call.Argument(0).ToString()
	if err != nil {
		return otto.Value{}, err
	}

	if !filepath.IsAbs(path) {
		dir := filepath.Dir(call.Otto.Context().Filename)
		path = filepath.Join(dir, path)
	}
	file, err := util.ReadFile(path)
	if err != nil {
		return otto.Value{}, err
	}

	return call.Otto.ToValue(file)
}

// Get returns the value contained at the given index
func (cs ConnectionSlice) Get(ii int) interface{} {
	return cs[ii]
}

// Len returns the number of items in the slice
func (cs ConnectionSlice) Len() int {
	return len(cs)
}

func stitchError(vm *otto.Otto, err error) otto.Value {
	return vm.MakeCustomError("StitchError", err.Error())
}

// toOttoFunc converts functions that return an error as a return value into
// a function that panics on errors. Otto requires functions to panic to signify
// errors in order to generate a stack trace.
func toOttoFunc(fn func(otto.FunctionCall) (otto.Value, error)) func(
	otto.FunctionCall) otto.Value {

	return func(call otto.FunctionCall) otto.Value {
		res, err := fn(call)
		if err != nil {
			// Otto uses `panic` with `*otto.Error`s to signify Javascript
			// runtime errors.
			if _, ok := err.(*otto.Error); ok {
				panic(err)
			}
			panic(stitchError(call.Otto, err))
		}
		return res
	}
}
<commit_msg>scheduler: Allow co-locating containers that share outbound ports<commit_after>//go:generate ../scripts/generate-bindings bindings.js

package stitch

import (
	"crypto/sha1"
	"encoding/json"
	"fmt"
	"path/filepath"

	"github.com/robertkrimen/otto"

	// Automatically import the Javascript underscore utility-belt library into
	// the Stitch VM.
	_ "github.com/robertkrimen/otto/underscore"

	"github.com/quilt/quilt/util"
)

// A Stitch is an abstract representation of the policy language.
type Stitch struct {
	Containers  []Container  `json:",omitempty"`
	Labels      []Label      `json:",omitempty"`
	Connections []Connection `json:",omitempty"`
	Placements  []Placement  `json:",omitempty"`
	Machines    []Machine    `json:",omitempty"`

	AdminACL  []string `json:",omitempty"`
	MaxPrice  float64  `json:",omitempty"`
	Namespace string   `json:",omitempty"`

	Invariants []invariant `json:",omitempty"`
}

// A Placement constraint guides where containers may be scheduled, either relative to
// the labels of other containers, or the machine the container will run on.
type Placement struct {
	TargetLabel string `json:",omitempty"`

	Exclusive bool `json:",omitempty"`

	// Label Constraint
	OtherLabel string `json:",omitempty"`

	// Machine Constraints
	Provider   string `json:",omitempty"`
	Size       string `json:",omitempty"`
	Region     string `json:",omitempty"`
	FloatingIP string `json:",omitempty"`
}

// An Image represents a Docker image that can be run. If the Dockerfile is non-empty,
// the image should be built and hosted by Quilt.
type Image struct {
	Name       string `json:",omitempty"`
	Dockerfile string `json:",omitempty"`
}

// A Container may be instantiated in the stitch and queried by users.
type Container struct {
	ID                string            `json:",omitempty"`
	Image             Image             `json:",omitempty"`
	Command           []string          `json:",omitempty"`
	Env               map[string]string `json:",omitempty"`
	FilepathToContent map[string]string `json:",omitempty"`
	Hostname          string            `json:",omitempty"`
}

// A Label represents a logical group of containers.
type Label struct {
	Name        string   `json:",omitempty"`
	IDs         []string `json:",omitempty"`
	Annotations []string `json:",omitempty"`
}

// A Connection allows containers implementing the From label to speak to containers
// implementing the To label in ports in the range [MinPort, MaxPort]
type Connection struct {
	From    string `json:",omitempty"`
	To      string `json:",omitempty"`
	MinPort int    `json:",omitempty"`
	MaxPort int    `json:",omitempty"`
}

// A ConnectionSlice allows for slices of Collections to be used in joins
type ConnectionSlice []Connection

// A Machine specifies the type of VM that should be booted.
type Machine struct {
	ID         string   `json:",omitempty"`
	Provider   string   `json:",omitempty"`
	Role       string   `json:",omitempty"`
	Size       string   `json:",omitempty"`
	CPU        Range    `json:",omitempty"`
	RAM        Range    `json:",omitempty"`
	DiskSize   int      `json:",omitempty"`
	Region     string   `json:",omitempty"`
	SSHKeys    []string `json:",omitempty"`
	FloatingIP string   `json:",omitempty"`
}

// A Range defines a range of acceptable values for a Machine attribute
type Range struct {
	Min float64 `json:",omitempty"`
	Max float64 `json:",omitempty"`
}

// PublicInternetLabel is a magic label that allows connections to or from the public
// network.
const PublicInternetLabel = "public"

// Accepts returns true if `x` is within the range specified by `stitchr` (include),
// or if no max is specified and `x` is larger than `stitchr.min`.
func (stitchr Range) Accepts(x float64) bool {
	return stitchr.Min <= x && (stitchr.Max == 0 || x <= stitchr.Max)
}

func run(vm *otto.Otto, filename string, code string) (otto.Value, error) {
	// Compile before running so that stacktraces have filenames.
	script, err := vm.Compile(filename, code)
	if err != nil {
		return otto.Value{}, err
	}

	return vm.Run(script)
}

func newVM(getter ImportGetter) (*otto.Otto, error) {
	vm := otto.New()
	if err := vm.Set("githubKeys", toOttoFunc(githubKeysImpl)); err != nil {
		return vm, err
	}
	if err := vm.Set("require", toOttoFunc(getter.requireImpl)); err != nil {
		return vm, err
	}
	if err := vm.Set("hash", toOttoFunc(hashImpl)); err != nil {
		return vm, err
	}
	if err := vm.Set("read", toOttoFunc(readImpl)); err != nil {
		return vm, err
	}

	_, err := run(vm, "<javascript_bindings>", javascriptBindings)
	return vm, err
}

// `runSpec` evaluates `spec` within a module closure.
func runSpec(vm *otto.Otto, filename string, spec string) (otto.Value, error) {
	// The function declaration must be prepended to the first line of the
	// import or else stacktraces will show an offset line number.
	exec := "(function() {" +
		"var module={exports: {}};" +
		"(function(module, exports) {" +
		spec +
		"})(module, module.exports);" +
		"return module.exports" +
		"})()"
	return run(vm, filename, exec)
}

// New parses and executes a stitch (in text form), and returns an abstract Dsl handle.
func New(filename string, specStr string, getter ImportGetter) (Stitch, error) {
	vm, err := newVM(getter)
	if err != nil {
		return Stitch{}, err
	}

	if _, err := runSpec(vm, filename, specStr); err != nil {
		return Stitch{}, err
	}

	spec, err := parseContext(vm)
	if err != nil {
		return Stitch{}, err
	}
	spec.createPortRules()

	if len(spec.Invariants) == 0 {
		return spec, nil
	}

	graph, err := InitializeGraph(spec)
	if err != nil {
		return Stitch{}, err
	}

	if err := checkInvariants(graph, spec.Invariants); err != nil {
		return Stitch{}, err
	}

	return spec, nil
}

// FromJavascript gets a Stitch handle from a string containing Javascript code.
func FromJavascript(specStr string, getter ImportGetter) (Stitch, error) {
	return New("<raw_string>", specStr, getter)
}

// FromFile gets a Stitch handle from a file on disk.
func FromFile(filename string, getter ImportGetter) (Stitch, error) {
	specStr, err := util.ReadFile(filename)
	if err != nil {
		return Stitch{}, err
	}
	return New(filename, specStr, getter)
}

// FromJSON gets a Stitch handle from the deployment representation.
func FromJSON(jsonStr string) (stc Stitch, err error) {
	err = json.Unmarshal([]byte(jsonStr), &stc)
	return stc, err
}

func parseContext(vm *otto.Otto) (stc Stitch, err error) {
	vmCtx, err := vm.Run("deployment.toQuiltRepresentation()")
	if err != nil {
		return stc, err
	}

	// Export() always returns `nil` as the error (it's only present for
	// backwards compatibility), so we can safely ignore it.
	exp, _ := vmCtx.Export()
	ctxStr, err := json.Marshal(exp)
	if err != nil {
		return stc, err
	}
	err = json.Unmarshal(ctxStr, &stc)
	return stc, err
}

// createPortRules creates exclusive placement rules such that no two containers
// listening on the same public port get placed on the same machine.
func (stitch *Stitch) createPortRules() {
	ports := make(map[int][]string)
	for _, c := range stitch.Connections {
		if c.From != PublicInternetLabel {
			continue
		}

		min := c.MinPort
		ports[min] = append(ports[min], c.To)
	}

	for _, labels := range ports {
		for _, tgt := range labels {
			for _, other := range labels {
				stitch.Placements = append(stitch.Placements,
					Placement{
						Exclusive:   true,
						TargetLabel: tgt,
						OtherLabel:  other,
					})
			}
		}
	}
}

// String returns the Stitch in its deployment representation.
func (stitch Stitch) String() string {
	jsonBytes, err := json.Marshal(stitch)
	if err != nil {
		panic(err)
	}
	return string(jsonBytes)
}

func hashImpl(call otto.FunctionCall) (otto.Value, error) {
	if len(call.ArgumentList) < 1 {
		panic(call.Otto.MakeRangeError(
			"hash requires an argument"))
	}

	toHash, err := call.Argument(0).ToString()
	if err != nil {
		return otto.Value{}, err
	}

	return call.Otto.ToValue(fmt.Sprintf("%x", sha1.Sum([]byte(toHash))))
}

func readImpl(call otto.FunctionCall) (otto.Value, error) {
	if len(call.ArgumentList) < 1 {
		panic(call.Otto.MakeRangeError(
			"read requires an argument"))
	}

	path, err := call.Argument(0).ToString()
	if err != nil {
		return otto.Value{}, err
	}

	if !filepath.IsAbs(path) {
		dir := filepath.Dir(call.Otto.Context().Filename)
		path = filepath.Join(dir, path)
	}
	file, err := util.ReadFile(path)
	if err != nil {
		return otto.Value{}, err
	}

	return call.Otto.ToValue(file)
}

// Get returns the value contained at the given index
func (cs ConnectionSlice) Get(ii int) interface{} {
	return cs[ii]
}

// Len returns the number of items in the slice
func (cs ConnectionSlice) Len() int {
	return len(cs)
}

func stitchError(vm *otto.Otto, err error) otto.Value {
	return vm.MakeCustomError("StitchError", err.Error())
}

// toOttoFunc converts functions that return an error as a return value into
// a function that panics on errors. Otto requires functions to panic to signify
// errors in order to generate a stack trace.
func toOttoFunc(fn func(otto.FunctionCall) (otto.Value, error)) func(
	otto.FunctionCall) otto.Value {

	return func(call otto.FunctionCall) otto.Value {
		res, err := fn(call)
		if err != nil {
			// Otto uses `panic` with `*otto.Error`s to signify Javascript
			// runtime errors.
			if _, ok := err.(*otto.Error); ok {
				panic(err)
			}
			panic(stitchError(call.Otto, err))
		}
		return res
	}
}
<|endoftext|>


<commit_before>package database

import (
	"context"
	"errors"
	"math/rand"
	"net/http"
	"time"

	"cloud.google.com/go/datastore"
	"github.com/Xe/hideyhole-site/interop"
	"github.com/bwmarrin/snowflake"
)

// Errors
var (
	ErrNoUserFound = errors.New("database: no user found")
	ErrNoFicFound  = errors.New("database: no fic found")
)

type Database struct {
	ds    *datastore.Client
	ctx   context.Context
	IDGen *snowflake.Node
}

//hack
func init() {
	http.DefaultServeMux = http.NewServeMux()
}

func Init(namespace, projectID string) (*Database, error) {
	ctx := context.Background()

	ctx = datastore.WithNamespace(ctx, namespace)

	client, err := datastore.NewClient(ctx, projectID)
	if err != nil {
		return nil, err
	}

	node, err := snowflake.NewNode(rand.Int63() % 1024)
	if err != nil {
		return nil, err
	}

	db := &Database{
		ds:    client,
		ctx:   ctx,
		IDGen: node,
	}

	return db, err
}

func (d *Database) GetUser(id string) (*interop.DiscordUser, error) {
	tx, err := d.ds.NewTransaction(d.ctx, datastore.MaxAttempts(1))
	if err != nil {
		return nil, err
	}
	defer tx.Rollback()

	dUser, _, err := d.getUser(d.ctx, tx, id)
	if err != nil {
		return nil, err
	}

	_, err = tx.Commit()
	if err != nil {
		return nil, err
	}

	return dUser, err
}

func (d *Database) getUser(ctx context.Context, tx *datastore.Transaction, id string) (*interop.DiscordUser, *datastore.Key, error) {
	results := []interop.DiscordUser{}
	var result *interop.DiscordUser
	var resultKey *datastore.Key

	q := datastore.NewQuery("DiscordUser").
		Filter("ID =", id).
		Limit(1).
		Transaction(tx)

	resultKeys, err := d.ds.GetAll(ctx, q, &results)
	if err != nil {
		switch err {
		case datastore.ErrNoSuchEntity:
			return nil, nil, ErrNoUserFound
		default:
			return nil, nil, err
		}
	}

	result = &results[0]
	resultKey = resultKeys[0]

	if len(result.ID) == 0 {
		return nil, nil, ErrNoUserFound
	}

	return result, resultKey, nil
}

func (d *Database) PutUser(dUser *interop.DiscordUser) error {
	tx, err := d.ds.NewTransaction(d.ctx, datastore.MaxAttempts(1))
	if err != nil {
		return err
	}
	defer tx.Rollback()

	_, err = d.putUser(d.ctx, tx, dUser)
	if err != nil {
		return err
	}

	_, err = tx.Commit()
	if err != nil {
		return err
	}

	return err
}

func (d *Database) putUser(ctx context.Context, tx *datastore.Transaction, dUser *interop.DiscordUser) (*datastore.PendingKey, error) {
	_, key, err := d.getUser(ctx, tx, dUser.ID)
	if err != nil && err != ErrNoUserFound {
		tx.Rollback()
		return nil, err
	}

	if key == nil {
		key = datastore.NewIncompleteKey(ctx, "DiscordUser", nil)
	}

	return tx.Put(key, dUser)
}

type Fic struct {
	ID         string
	Created    time.Time
	Edited     time.Time
	AuthorID   string
	Title      string
	ChapterIDs []string

	Description     string `datastore:",noindex"`
	DescriptionHTML string `datastore:",noindex"`
}

func (d *Database) GetFic(id string) (*Fic, error) {
	tx, err := d.ds.NewTransaction(d.ctx, datastore.MaxAttempts(1))
	if err != nil {
		return nil, err
	}
	defer tx.Rollback()

	fic, _, err := d.getFic(d.ctx, tx, id)
	if err != nil {
		return nil, err
	}

	_, err = tx.Commit()
	if err != nil {
		return nil, err
	}

	return fic, nil
}

func (d *Database) getFic(ctx context.Context, tx *datastore.Transaction, id string) (*Fic, *datastore.Key, error) {
	results := []Fic{}
	var result *Fic
	var resultKey *datastore.Key

	q := datastore.NewQuery("Fic").
		Filter("ID =", id).
		Limit(1).
		Transaction(tx)

	resultKeys, err := d.ds.GetAll(ctx, q, results)
	if err != nil {
		switch err {
		case datastore.ErrNoSuchEntity:
			return nil, nil, ErrNoFicFound
		default:
			return nil, nil, err
		}
	}

	result = &results[0]
	resultKey = resultKeys[0]

	return result, resultKey, nil
}

func (d *Database) putFic(ctx context.Context, tx *datastore.Transaction, fic *Fic) (*datastore.PendingKey, error) {
	_, key, err := d.getFic(ctx, tx, fic.ID)
	if err != nil && err != ErrNoFicFound {
		return nil, err
	}

	if key == nil {
		key = datastore.NewIncompleteKey(ctx, "Fic", nil)
	}

	return tx.Put(key, fic)
}

type Chapter struct {
	ID       string
	FicID    string
	AuthorID string
	Title    string
	Created  time.Time

	Content     string `datastore:",noindex"`
	ContentHTML string `datastore:",noindex"`
}

func (d *Database) addFicChapter(ctx context.Context, tx *datastore.Transaction, fic *Fic, ficKey *datastore.Key, chapter *Chapter) (*datastore.PendingKey, error) {
	fic.ChapterIDs = append(fic.ChapterIDs, chapter.ID)

	_, err := d.putFic(ctx, tx, fic)
	if err != nil {
		return nil, err
	}

	key := datastore.NewIncompleteKey(ctx, "Chapter", ficKey)

	return tx.Put(key, chapter)
}
<commit_msg>database: query fic and chapters<commit_after>package database

import (
	"context"
	"errors"
	"math/rand"
	"net/http"
	"time"

	"cloud.google.com/go/datastore"
	"github.com/Xe/hideyhole-site/interop"
	"github.com/bwmarrin/snowflake"
)

// Errors
var (
	ErrNoUserFound = errors.New("database: no user found")
	ErrNoFicFound  = errors.New("database: no fic found")
)

type Database struct {
	ds    *datastore.Client
	ctx   context.Context
	IDGen *snowflake.Node
}

//hack
func init() {
	http.DefaultServeMux = http.NewServeMux()
}

func Init(namespace, projectID string) (*Database, error) {
	ctx := context.Background()

	ctx = datastore.WithNamespace(ctx, namespace)

	client, err := datastore.NewClient(ctx, projectID)
	if err != nil {
		return nil, err
	}

	node, err := snowflake.NewNode(rand.Int63() % 1024)
	if err != nil {
		return nil, err
	}

	db := &Database{
		ds:    client,
		ctx:   ctx,
		IDGen: node,
	}

	return db, err
}

func (d *Database) GetUser(id string) (*interop.DiscordUser, error) {
	tx, err := d.ds.NewTransaction(d.ctx, datastore.MaxAttempts(1))
	if err != nil {
		return nil, err
	}
	defer tx.Rollback()

	dUser, _, err := d.getUser(d.ctx, tx, id)
	if err != nil {
		return nil, err
	}

	_, err = tx.Commit()
	if err != nil {
		return nil, err
	}

	return dUser, err
}

func (d *Database) getUser(ctx context.Context, tx *datastore.Transaction, id string) (*interop.DiscordUser, *datastore.Key, error) {
	results := []interop.DiscordUser{}
	var result *interop.DiscordUser
	var resultKey *datastore.Key

	q := datastore.NewQuery("DiscordUser").
		Filter("ID =", id).
		Limit(1).
		Transaction(tx)

	resultKeys, err := d.ds.GetAll(ctx, q, &results)
	if err != nil {
		switch err {
		case datastore.ErrNoSuchEntity:
			return nil, nil, ErrNoUserFound
		default:
			return nil, nil, err
		}
	}

	result = &results[0]
	resultKey = resultKeys[0]

	if len(result.ID) == 0 {
		return nil, nil, ErrNoUserFound
	}

	return result, resultKey, nil
}

func (d *Database) PutUser(dUser *interop.DiscordUser) error {
	tx, err := d.ds.NewTransaction(d.ctx, datastore.MaxAttempts(1))
	if err != nil {
		return err
	}
	defer tx.Rollback()

	_, err = d.putUser(d.ctx, tx, dUser)
	if err != nil {
		return err
	}

	_, err = tx.Commit()
	if err != nil {
		return err
	}

	return err
}

func (d *Database) putUser(ctx context.Context, tx *datastore.Transaction, dUser *interop.DiscordUser) (*datastore.PendingKey, error) {
	_, key, err := d.getUser(ctx, tx, dUser.ID)
	if err != nil && err != ErrNoUserFound {
		tx.Rollback()
		return nil, err
	}

	if key == nil {
		key = datastore.NewIncompleteKey(ctx, "DiscordUser", nil)
	}

	return tx.Put(key, dUser)
}

type Fic struct {
	ID         string
	Created    time.Time
	Edited     time.Time
	AuthorID   string
	Title      string
	ChapterIDs []string

	Description     string `datastore:",noindex"`
	DescriptionHTML string `datastore:",noindex"`
}

func (d *Database) GetFic(id string) (*Fic, error) {
	tx, err := d.ds.NewTransaction(d.ctx, datastore.MaxAttempts(1))
	if err != nil {
		return nil, err
	}
	defer tx.Rollback()

	fic, _, err := d.getFic(d.ctx, tx, id)
	if err != nil {
		return nil, err
	}

	_, err = tx.Commit()
	if err != nil {
		return nil, err
	}

	return fic, nil
}

func (d *Database) getFic(ctx context.Context, tx *datastore.Transaction, id string) (*Fic, *datastore.Key, error) {
	results := []Fic{}
	var result *Fic
	var resultKey *datastore.Key

	q := datastore.NewQuery("Fic").
		Filter("ID =", id).
		Limit(1).
		Transaction(tx)

	resultKeys, err := d.ds.GetAll(ctx, q, results)
	if err != nil {
		switch err {
		case datastore.ErrNoSuchEntity:
			return nil, nil, ErrNoFicFound
		default:
			return nil, nil, err
		}
	}

	result = &results[0]
	resultKey = resultKeys[0]

	return result, resultKey, nil
}

func (d *Database) GetFicAndChapters(ficID string) (*Fic, []Chapter, error) {
	tx, err := d.ds.NewTransaction(d.ctx, datastore.MaxAttempts(1))
	if err != nil {
		return nil, nil, err
	}
	defer tx.Rollback()

	fic, _, err := d.getFic(d.ctx, tx, ficID)
	if err != nil {
		return nil, nil, err
	}

	q := datastore.NewQuery("Chapter").
		Filter("FicID =", ficID).
		Transaction(tx)

	var chapters []Chapter
	_, err = d.ds.GetAll(d.ctx, q, &chapters)
	if err != nil {
		return nil, nil, err
	}

	return fic, chapters, nil
}

func (d *Database) putFic(ctx context.Context, tx *datastore.Transaction, fic *Fic) (*datastore.PendingKey, error) {
	_, key, err := d.getFic(ctx, tx, fic.ID)
	if err != nil && err != ErrNoFicFound {
		return nil, err
	}

	if key == nil {
		key = datastore.NewIncompleteKey(ctx, "Fic", nil)
	}

	return tx.Put(key, fic)
}

type Chapter struct {
	ID       string
	FicID    string
	AuthorID string
	Title    string
	Created  time.Time

	Content     string `datastore:",noindex"`
	ContentHTML string `datastore:",noindex"`
}

func (d *Database) addFicChapter(ctx context.Context, tx *datastore.Transaction, fic *Fic, ficKey *datastore.Key, chapter *Chapter) (*datastore.PendingKey, error) {
	fic.ChapterIDs = append(fic.ChapterIDs, chapter.ID)

	_, err := d.putFic(ctx, tx, fic)
	if err != nil {
		return nil, err
	}

	key := datastore.NewIncompleteKey(ctx, "Chapter", ficKey)

	return tx.Put(key, chapter)
}
<|endoftext|>


<commit_before>package storage

import (
	"time"

	cache "github.com/patrickmn/go-cache"
)

type Cache interface {
	Get(string) ([]byte, bool)
	Set(int, string, []byte)
	Clear()
}

type dbOperation int

const (
	dbPut = iota
	dbDelete
)

const (
	defaultTimeout    = 1 * time.Minute
	defaultExpiration = 2 * time.Minute
)

type dbCache struct {
	cache *cache.Cache
}

type cacheData struct {
	op    int
	value []byte
}

func newCache() Cache {
	return &dbCache{
		cache: cache.New(defaultTimeout, defaultExpiration),
	}
}

func (c *dbCache) Get(key string) ([]byte, bool) {
	obj, found := c.cache.Get(key)
	if !found {
		return []byte{}, found
	}

	data := obj.(cacheData)
	// if key is deleted, then cache should return not found
	if dbDelete == data.op {
		return []byte{}, false
	}

	return data.value, found
}

func (c *dbCache) Set(op int, key string, value []byte) {
	cached := cacheData{
		op:    op,
		value: value,
	}
	c.cache.Set(key, cached, defaultExpiration)
}

func (c *dbCache) Clear() {
	c.cache.Flush()
}
<commit_msg>[storage] - Fixes type dbOperation is unused (U1000)<commit_after>package storage

import (
	"time"

	cache "github.com/patrickmn/go-cache"
)

type Cache interface {
	Get(string) ([]byte, bool)
	Set(int, string, []byte)
	Clear()
}

const (
	dbPut = iota
	dbDelete
)

const (
	defaultTimeout    = 1 * time.Minute
	defaultExpiration = 2 * time.Minute
)

type dbCache struct {
	cache *cache.Cache
}

type cacheData struct {
	op    int
	value []byte
}

func newCache() Cache {
	return &dbCache{
		cache: cache.New(defaultTimeout, defaultExpiration),
	}
}

func (c *dbCache) Get(key string) ([]byte, bool) {
	obj, found := c.cache.Get(key)
	if !found {
		return []byte{}, found
	}

	data := obj.(cacheData)
	// if key is deleted, then cache should return not found
	if dbDelete == data.op {
		return []byte{}, false
	}

	return data.value, found
}

func (c *dbCache) Set(op int, key string, value []byte) {
	cached := cacheData{
		op:    op,
		value: value,
	}
	c.cache.Set(key, cached, defaultExpiration)
}

func (c *dbCache) Clear() {
	c.cache.Flush()
}
<|endoftext|>


<commit_before>package str

import (
	"strings"
	"unicode"

	"golang.org/x/text/runes"

	"golang.org/x/text/transform"
	"golang.org/x/text/unicode/norm"
)

//RemoveWhiteSpace removes whitespace {'\n' '\t' ' ' '\r`} from a string.
//Note that this converts to runes and back to UTF-8, so RemoveWhiteSpace(s) == s
//for a non-whitespace string does not necessarially hold, since the code points may differ.
func RemoveWhiteSpace(s string) string {
	return RemoveRunes(s, '\n', ' ', '\r', '\t')
}

//NormEqual returns true if the NKFC normalized forms of both strings are equal.
func NormEqual(s, q string) bool {
	return NKFC(s) == NKFC(q)
}

//NormFoldEqual returns true if the casefolded, NKFC normalized forms of both strings are equal.
func NormFoldEqual(s, q string) bool {
	return strings.EqualFold(NKFC(s), NKFC(q))
}

//NKFC normalizes a string to it's NKFC form
func NKFC(s string) string {
	return norm.NFKC.String(s)
}

//NFKD normalizes a string to it's NFKD form
func NFKD(s string) string {
	return norm.NFKD.String(s)
}

//NFD normalizes a string to it's NFD form
func NFD(s string) string {
	return norm.NFD.String(s)
}

//NFC normalizes a string to it's NFC form
func NFC(s string) string {
	return norm.NFC.String(s)
}

func isNonSpacingMark(r rune) bool {
	return unicode.Is(unicode.Mn, r) // Mn: nonspacing marks
}

var removeNonSpacingMarks = runes.Remove(runes.In(unicode.Mn))
var diacriticRemover = transform.Chain(norm.NFD, removeNonSpacingMarks, norm.NFC)

//RemoveDiacriticsNFC creates a copy of s with the diacritics removed. It also transforms it to NFC.
//Thread Safe
func RemoveDiacriticsNFC(s string) string {
	var diacriticRemover = transform.Chain(norm.NFD, removeNonSpacingMarks, norm.NFC)
	out, _, _ := transform.String(diacriticRemover, s)
	return out
}
<commit_msg>added more normalization options<commit_after>package str

import (
	"strings"
	"unicode"

	"golang.org/x/text/runes"

	"golang.org/x/text/transform"
	"golang.org/x/text/unicode/norm"
)

//NormEqual returns true if the NKFC normalized forms of both strings are equal.
func NormEqual(s, q string) bool {
	return NKFC(s) == NKFC(q)
}

//NormFoldEqual returns true if the casefolded, NKFC normalized forms of both strings are equal.
func NormFoldEqual(s, q string) bool {
	return strings.EqualFold(NKFC(s), NKFC(q))
}

//NKFC normalizes a string to it's NKFC form
func NKFC(s string) string {
	return norm.NFKC.String(s)
}

//NFKD normalizes a string to it's NFKD form
func NFKD(s string) string {
	return norm.NFKD.String(s)
}

//NFD normalizes a string to it's NFD form
func NFD(s string) string {
	return norm.NFD.String(s)
}

//NFC normalizes a string to it's NFC form
func NFC(s string) string {
	return norm.NFC.String(s)
}

func isNonSpacingMark(r rune) bool {
	return unicode.Is(unicode.Mn, r) // Mn: nonspacing marks
}

var (
	removeNonSpacingMarks = runes.Remove(runes.In(unicode.Mn))
	removeWhiteSpace      = runes.Remove(runes.In(unicode.White_Space))
	removePunctuation     = runes.Remove(runes.In(unicode.Punct))
	removeControl         = runes.Remove(runes.In(unicode.C))
)

//RemoveDiacriticsNFC creates a copy of s with the diacritics removed. It also transforms it to NFC.
//Thread Safe
func RemoveDiacriticsNFC(s string) string {
	var diacriticRemover = transform.Chain(norm.NFD, removeNonSpacingMarks, norm.NFC)
	out, _, _ := transform.String(diacriticRemover, s)
	return out
}

//RemovePunctuationNTS removes punctuation (as defined by unicode) from a string, but is NOT thread safe.
//Note that this converts to runes and back to UTF-8, so RemoveWhiteSpace(s) == s
//for a string that contains non-punctuation characters does not necessarially hold, since the code points may differ.
func RemovePunctuationNTS(s string) string {
	out, _, _ := transform.String(removePunctuation, s)
	return out
}

//RemoveWhiteSpaceNTS removes whitespace (as defined by unicode) from a string, but is NOT thread safe.
//Note that this converts to runes and back to UTF-8, so RemoveWhiteSpace(s) == s
//for a non-whitespace string does not necessarially hold, since the code points may differ.
//Note that this is faster than RemoveWhitespace, but is not thread safe.
func RemoveWhiteSpaceNTS(s string) string {
	out, _, _ := transform.String(removePunctuation, s)
	return out
}

//RemoveControlNTS removes whitespace (as defined by unicode) from a string, but is NOT thread safe.
//Note that this converts to runes and back to UTF-8, so RemoveWhiteSpace(s) == s
//for a non-whitespace string does not necessarially hold, since the code points may differ.
//Note that this is faster than RemoveWhitespace, but is not thread safe.
func RemoveControlNTS(s string) string {
	out, _, _ := transform.String(removeControl, s)
	return out
}

//ExtremeNormalization heavily normalizes a stirng for purposes of comparison and safety.
//It removes ALL nonspacing marks, whitespace, punctuation, control characters, and transforms the
//string to NFKC encoding. This can and will lose a lot of information!
func ExtremeNormalization(s string) string {
	s = strings.ToLower(s)
	normalizer := transform.Chain(
		norm.NFKD,
		removeNonSpacingMarks,
		removeWhiteSpace,
		removePunctuation,
		removeControl,
		norm.NFKC)
	out, _, _ := transform.String(normalizer, s)
	return out
}
<|endoftext|>