<commit_before>package main

import (
  "path"
  "os"
  "sync"
  "log"
  "sort"
  "io"
  "io/ioutil"
  "errors"
  "path/filepath"
)

const (
  CopyPrice = 100
  RenamePrice = 10

  RemoveFactor = RenamePrice
  UpdateFactor = RenamePrice + CopyPrice
  AddFactor = CopyPrice
)

type BackupPair struct {
  relpath string
  newpath string
}

type ProgressHandler interface {
  HandleSystemMessage(message string)
  HandlePercentChange(percent int)
  HandleFinish()
}

type LogProgressHandler struct {
}

type ProgressReporter struct {
  grandTotal uint64
  currentProgress uint64
  progressChan chan int64
  percent int //0..100
  reportingChan chan bool
  systemMessageChan chan string
  finished chan bool
  progressHandler ProgressHandler
}

type PackageInstaller struct {
  backups map[string]string
  backupsChan chan BackupPair
  progressReporter *ProgressReporter
  installDir string
  packageDir string
  failInTheEnd bool // for debugging purposes
}

func (pi *PackageInstaller) Install(filesProvider UpdateFilesProvider) error {
  pi.progressReporter.grandTotal = pi.calculateGrandTotals(filesProvider)
  go pi.progressReporter.reportingLoop()
  defer close(pi.progressReporter.progressChan)
  defer func() {
    go func () {
      pi.progressReporter.finished <- true
    }()
  }()

  err := pi.installPackage(filesProvider)

  if err == nil {
    pi.afterSuccess()
  } else {
    pi.afterFailure(filesProvider)
  }

  return err
}

func (pi *PackageInstaller) calculateGrandTotals(filesProvider UpdateFilesProvider) uint64 {
  var sum uint64

  for _, fi := range filesProvider.FilesToRemove() {
    sum += uint64(fi.FileSize * RemoveFactor) / 100
  }

  for _, fi := range filesProvider.FilesToUpdate() {
    sum += uint64(fi.FileSize * UpdateFactor) / 100
  }

  for _, fi := range filesProvider.FilesToAdd() {
    sum += uint64(fi.FileSize * AddFactor) / 100
  }

  return sum
}

func (pi *PackageInstaller) installPackage(filesProvider UpdateFilesProvider) (err error) {
  log.Println("Installing package...")

  var wg sync.WaitGroup
  wg.Add(1)
  go func() {
    for bp := range pi.backupsChan {
      pi.backups[bp.relpath] = bp.newpath
    }
    wg.Done()
  }()

  pi.progressReporter.systemMessageChan <- "Removing components"
  err = pi.removeFiles(filesProvider.FilesToRemove())
  if err != nil {
    return err
  }

  pi.progressReporter.systemMessageChan <- "Updating components"
  err = pi.updateFiles(filesProvider.FilesToUpdate())
  if err != nil {
    return err
  }

  pi.progressReporter.systemMessageChan <- "Adding components"
  err = pi.addFiles(filesProvider.FilesToAdd())
  if err != nil {
    return err
  }

  go func() {
    close(pi.backupsChan)
  }()

  wg.Wait()

  if pi.failInTheEnd {
    err = errors.New("Fail by demand")
  }

  return err
}

func (pi *PackageInstaller) afterSuccess() {
  log.Println("After success")
  pi.progressReporter.systemMessageChan <- "Finishing the installation..."
  cleanupEmptyDirs(pi.installDir)
  pi.removeBackups();
}

func (pi *PackageInstaller) afterFailure(filesProvider UpdateFilesProvider) {
  log.Println("After failure")
  pi.progressReporter.systemMessageChan <- "Cleaning up..."
  purgeFiles(pi.installDir, filesProvider.FilesToAdd())
  pi.restoreBackups()
  cleanupEmptyDirs(pi.installDir)
  pi.removeBackups()
}

func copyFile(src, dst string) (err error) {
  in, err := os.Open(src)
  if err != nil {
    log.Printf("Failed to open source: %v", err)
    return
  }

  defer in.Close()

  out, err := os.Create(dst)
  if err != nil {
    log.Printf("Failed to create destination: %v", err)
    return
  }

  defer func() {
    cerr := out.Close()
    if err == nil {
      err = cerr
    }
  }()

  if _, err = io.Copy(out, in); err != nil {
    return
  }

  err = out.Sync()
  return
}

func (pi *PackageInstaller) backupFile(relpath string) error {
  log.Printf("Backing up %v", relpath)

  oldpath := path.Join(pi.installDir, relpath)
  backupPath := relpath + ".bak"

  newpath := path.Join(pi.installDir, backupPath)

  err := os.Rename(oldpath, newpath)

  if err == nil {
    pi.backupsChan <- BackupPair{relpath: relpath, newpath: newpath}
  } else {
    log.Printf("Backup failed: %v", err)
  }

  return err
}

func (pi *PackageInstaller) restoreBackups() {
  log.Printf("Restoring %v backups", len(pi.backups))

  var wg sync.WaitGroup

  for relpath, backuppath := range pi.backups {
    wg.Add(1)

    relativePath := relpath
    pathToRestore := backuppath

    go func() {
      defer wg.Done()

      oldpath := path.Join(pi.installDir, relativePath)
      err := os.Rename(pathToRestore, oldpath)

      if err != nil {
        log.Println(err)
      }
    }()
  }

  wg.Wait()
}

func (pi *PackageInstaller) removeBackups() {
  log.Printf("Removing %v backups", len(pi.backups))

  var wg sync.WaitGroup

  for _, backuppath := range pi.backups {
    wg.Add(1)

    pathToRemove := backuppath

    go func() {
      defer wg.Done()

      err := os.Remove(pathToRemove)
      if err != nil {
        log.Println(err)
      }
    }()
  }

  wg.Wait()
}

func (pi *PackageInstaller) removeFiles(files []*UpdateFileInfo) error {
  log.Printf("Removing %v files", len(files))

  var wg sync.WaitGroup
  errc := make(chan error)
  done := make(chan bool)

  for _, fi := range files {
    wg.Add(1)
    pathToRemove, filesize := fi.Filepath, fi.FileSize

    go func() {
      defer wg.Done()

      select {
      case <-done: return
      default:
      }

      fullpath := filepath.Join(pi.installDir, pathToRemove)
      log.Printf("Removing file %v", fullpath)

      err := pi.backupFile(pathToRemove)

      if err != nil {
        log.Printf("Removing file %v failed", pathToRemove)
        log.Println(err)
        errc <- err
        close(done)
      } else {
        go pi.progressReporter.accountRemove(filesize)
      }
    }()
  }

  go func() {
    errc <- nil
  }()

  wg.Wait()

  if err := <-errc; err != nil {
    return err
  }

  return nil
}

func (pi *PackageInstaller) updateFiles(files []*UpdateFileInfo) error {
  log.Printf("Updating %v files", len(files))

  var wg sync.WaitGroup
  errc := make(chan error)
  done := make(chan bool)

  for _, fi := range files {
    wg.Add(1)

    pathToUpdate, filesize := fi.Filepath, fi.FileSize

    go func() {
      defer wg.Done()

      select {
      case <-done: return
      default:
      }

      oldpath := path.Join(pi.installDir, pathToUpdate)
      log.Printf("Updating file %v", oldpath)

      err := pi.backupFile(pathToUpdate)

      if err == nil {
        newpath := path.Join(pi.packageDir, pathToUpdate)
        err = os.Rename(newpath, oldpath)
      }

      if err != nil {
        log.Printf("Updating file %v failed", pathToUpdate)
        log.Println(err)
        errc <- err
        close(done)
      } else {
        go pi.progressReporter.accountUpdate(filesize)
      }
    }()
  }

  go func() {
    errc <- nil
  }()

  wg.Wait()

  if err := <-errc; err != nil {
    return err
  }

  return nil
}

func (pi *PackageInstaller) addFiles(files []*UpdateFileInfo) error {
  log.Printf("Adding %v files", len(files))

  var wg sync.WaitGroup
  errc := make(chan error)
  done := make(chan bool)

  for _, fi := range files {
    wg.Add(1)

    pathToAdd, filesize := fi.Filepath, fi.FileSize

    go func() {
      defer wg.Done()

      select {
      case <-done: return
      default:
      }

      oldpath := path.Join(pi.installDir, pathToAdd)
      ensureDirExists(oldpath)

      newpath := path.Join(pi.packageDir, pathToAdd)
      err := os.Rename(newpath, oldpath)

      if err != nil {
        log.Printf("Adding file %v failed", pathToAdd)
        log.Println(err)
        errc <- err
        close(done)
      } else {
        go pi.progressReporter.accountAdd(filesize)
      }
    }()
  }

  go func() {
    errc <- nil
  }()

  wg.Wait()

  if err := <-errc; err != nil {
    return err
  }

  return nil
}

func purgeFiles(root string, files []*UpdateFileInfo) {
  log.Printf("Purging %v files", len(files))

  var wg sync.WaitGroup

  for _, fi := range files {
    wg.Add(1)

    fileToPurge := fi.Filepath

    go func() {
      defer wg.Done()

      fullpath := path.Join(root, fileToPurge)
      err := os.Remove(fullpath)
      if err != nil {
        log.Println(err)
      }
    }()
  }

  wg.Wait()
}

func ensureDirExists(fullpath string) (err error) {
  dirpath := path.Dir(fullpath)
  err = os.MkdirAll(dirpath, os.ModeDir)
  if err != nil {
    log.Printf("Failed to create directory %v", dirpath)
  }

  return err
}

type ByLength []string

func (s ByLength) Len() int {
    return len(s)
}
func (s ByLength) Swap(i, j int) {
    s[i], s[j] = s[j], s[i]
}
func (s ByLength) Less(i, j int) bool {
    return len(s[i]) > len(s[j])
}

func cleanupEmptyDirs(root string) {
  c := make(chan string)

  go func() {
    var wg sync.WaitGroup
    err := filepath.Walk(root, func(path string, info os.FileInfo, err error) error {
      if err != nil {
        return err
      }

      if info.Mode().IsDir() {
        wg.Add(1)
        go func() {
          c <- path
          wg.Done()
        }()
      }

      return nil
    })

    if err != nil {
      log.Println(err)
    }

    go func() {
      wg.Wait()
      close(c)
    }()
  }()

  dirs := make([]string, 0)
  for path := range c {
    dirs = append(dirs, path)
  }

  removeEmptyDirs(dirs)
}

func removeEmptyDirs(dirs []string) {
  sort.Sort(ByLength(dirs))

  for _, dirpath := range dirs {
    entries, err := ioutil.ReadDir(dirpath)
    if err != nil { continue }

    if len(entries) == 0 {
      log.Printf("Removing empty dir %v", dirpath)

      err = os.Remove(dirpath)
      if err != nil {
        log.Println(err)
      }
    }
  }
}

func (pr *ProgressReporter) accountRemove(progress int64) {
  pr.progressChan <- (progress*RemoveFactor)/100
}

func (pr *ProgressReporter) accountUpdate(progress int64) {
  pr.progressChan <- (progress*UpdateFactor)/100
}

func (pr *ProgressReporter) accountAdd(progress int64) {
  pr.progressChan <- (progress*AddFactor)/100
}

func (pr *ProgressReporter) reportingLoop() {
  for chunk := range pr.progressChan {
    pr.currentProgress += uint64(chunk)

    percent := (pr.currentProgress*100) / pr.grandTotal
    pr.percent = int(percent)

    go func() {
      pr.reportingChan <- true
    }()
  }

  close(pr.reportingChan)
}

func (pr *ProgressReporter) receiveUpdates() {
  for _ = range pr.reportingChan {
    pr.progressHandler.HandlePercentChange(pr.percent)
  }
}

func (pr *ProgressReporter) receiveSystemMessages() {
  for msg := range pr.systemMessageChan {
    pr.progressHandler.HandleSystemMessage(msg)
  }
}

func (pr *ProgressReporter) receiveFinish() {
  <- pr.finished
  pr.progressHandler.HandleFinish()
}

func (pr *ProgressReporter) handleProgress() {
  go pr.receiveSystemMessages()
  go pr.receiveUpdates()
  go pr.receiveFinish()
}

func (ph *LogProgressHandler) HandlePercentChange(percent int) {
  log.Printf("Completed %v%%", percent)
}

func (ph *LogProgressHandler) HandleSystemMessage(msg string) {
  log.Printf("System message: %v", msg)
}

func (ph *LogProgressHandler) HandleFinish() {
  log.Printf("Finished")
}
<commit_msg>Fix for empty dirs<commit_after>package main

import (
  "path"
  "os"
  "sync"
  "log"
  "sort"
  "io"
  "io/ioutil"
  "errors"
  "path/filepath"
)

const (
  CopyPrice = 100
  RenamePrice = 10

  RemoveFactor = RenamePrice
  UpdateFactor = RenamePrice + CopyPrice
  AddFactor = CopyPrice
)

type BackupPair struct {
  relpath string
  newpath string
}

type ProgressHandler interface {
  HandleSystemMessage(message string)
  HandlePercentChange(percent int)
  HandleFinish()
}

type LogProgressHandler struct {
}

type ProgressReporter struct {
  grandTotal uint64
  currentProgress uint64
  progressChan chan int64
  percent int //0..100
  reportingChan chan bool
  systemMessageChan chan string
  finished chan bool
  progressHandler ProgressHandler
}

type PackageInstaller struct {
  backups map[string]string
  backupsChan chan BackupPair
  progressReporter *ProgressReporter
  installDir string
  packageDir string
  failInTheEnd bool // for debugging purposes
}

func (pi *PackageInstaller) Install(filesProvider UpdateFilesProvider) error {
  pi.progressReporter.grandTotal = pi.calculateGrandTotals(filesProvider)
  go pi.progressReporter.reportingLoop()
  defer close(pi.progressReporter.progressChan)
  defer func() {
    go func () {
      pi.progressReporter.finished <- true
    }()
  }()

  err := pi.installPackage(filesProvider)

  if err == nil {
    pi.afterSuccess()
  } else {
    pi.afterFailure(filesProvider)
  }

  return err
}

func (pi *PackageInstaller) calculateGrandTotals(filesProvider UpdateFilesProvider) uint64 {
  var sum uint64

  for _, fi := range filesProvider.FilesToRemove() {
    sum += uint64(fi.FileSize * RemoveFactor) / 100
  }

  for _, fi := range filesProvider.FilesToUpdate() {
    sum += uint64(fi.FileSize * UpdateFactor) / 100
  }

  for _, fi := range filesProvider.FilesToAdd() {
    sum += uint64(fi.FileSize * AddFactor) / 100
  }

  return sum
}

func (pi *PackageInstaller) installPackage(filesProvider UpdateFilesProvider) (err error) {
  log.Println("Installing package...")

  var wg sync.WaitGroup
  wg.Add(1)
  go func() {
    for bp := range pi.backupsChan {
      pi.backups[bp.relpath] = bp.newpath
    }
    wg.Done()
  }()

  pi.progressReporter.systemMessageChan <- "Removing components"
  err = pi.removeFiles(filesProvider.FilesToRemove())
  if err != nil {
    return err
  }

  pi.progressReporter.systemMessageChan <- "Updating components"
  err = pi.updateFiles(filesProvider.FilesToUpdate())
  if err != nil {
    return err
  }

  pi.progressReporter.systemMessageChan <- "Adding components"
  err = pi.addFiles(filesProvider.FilesToAdd())
  if err != nil {
    return err
  }

  go func() {
    close(pi.backupsChan)
  }()

  wg.Wait()

  if pi.failInTheEnd {
    err = errors.New("Fail by demand")
  }

  return err
}

func (pi *PackageInstaller) afterSuccess() {
  log.Println("After success")
  pi.progressReporter.systemMessageChan <- "Finishing the installation..."
  pi.removeBackups();
  cleanupEmptyDirs(pi.installDir)
}

func (pi *PackageInstaller) afterFailure(filesProvider UpdateFilesProvider) {
  log.Println("After failure")
  pi.progressReporter.systemMessageChan <- "Cleaning up..."
  purgeFiles(pi.installDir, filesProvider.FilesToAdd())
  pi.restoreBackups()
  pi.removeBackups()
  cleanupEmptyDirs(pi.installDir)
}

func copyFile(src, dst string) (err error) {
  in, err := os.Open(src)
  if err != nil {
    log.Printf("Failed to open source: %v", err)
    return
  }

  defer in.Close()

  out, err := os.Create(dst)
  if err != nil {
    log.Printf("Failed to create destination: %v", err)
    return
  }

  defer func() {
    cerr := out.Close()
    if err == nil {
      err = cerr
    }
  }()

  if _, err = io.Copy(out, in); err != nil {
    return
  }

  err = out.Sync()
  return
}

func (pi *PackageInstaller) backupFile(relpath string) error {
  log.Printf("Backing up %v", relpath)

  oldpath := path.Join(pi.installDir, relpath)
  backupPath := relpath + ".bak"

  newpath := path.Join(pi.installDir, backupPath)

  err := os.Rename(oldpath, newpath)

  if err == nil {
    pi.backupsChan <- BackupPair{relpath: relpath, newpath: newpath}
  } else {
    log.Printf("Backup failed: %v", err)
  }

  return err
}

func (pi *PackageInstaller) restoreBackups() {
  log.Printf("Restoring %v backups", len(pi.backups))

  var wg sync.WaitGroup

  for relpath, backuppath := range pi.backups {
    wg.Add(1)

    relativePath := relpath
    pathToRestore := backuppath

    go func() {
      defer wg.Done()

      oldpath := path.Join(pi.installDir, relativePath)
      err := os.Rename(pathToRestore, oldpath)

      if err != nil {
        log.Println(err)
      }
    }()
  }

  wg.Wait()
}

func (pi *PackageInstaller) removeBackups() {
  log.Printf("Removing %v backups", len(pi.backups))

  var wg sync.WaitGroup

  for _, backuppath := range pi.backups {
    wg.Add(1)

    pathToRemove := backuppath

    go func() {
      defer wg.Done()

      err := os.Remove(pathToRemove)
      if err != nil {
        log.Println(err)
      }
    }()
  }

  wg.Wait()
}

func (pi *PackageInstaller) removeFiles(files []*UpdateFileInfo) error {
  log.Printf("Removing %v files", len(files))

  var wg sync.WaitGroup
  errc := make(chan error)
  done := make(chan bool)

  for _, fi := range files {
    wg.Add(1)
    pathToRemove, filesize := fi.Filepath, fi.FileSize

    go func() {
      defer wg.Done()

      select {
      case <-done: return
      default:
      }

      fullpath := filepath.Join(pi.installDir, pathToRemove)
      log.Printf("Removing file %v", fullpath)

      err := pi.backupFile(pathToRemove)

      if err != nil {
        log.Printf("Removing file %v failed", pathToRemove)
        log.Println(err)
        errc <- err
        close(done)
      } else {
        go pi.progressReporter.accountRemove(filesize)
      }
    }()
  }

  go func() {
    errc <- nil
  }()

  wg.Wait()

  if err := <-errc; err != nil {
    return err
  }

  return nil
}

func (pi *PackageInstaller) updateFiles(files []*UpdateFileInfo) error {
  log.Printf("Updating %v files", len(files))

  var wg sync.WaitGroup
  errc := make(chan error)
  done := make(chan bool)

  for _, fi := range files {
    wg.Add(1)

    pathToUpdate, filesize := fi.Filepath, fi.FileSize

    go func() {
      defer wg.Done()

      select {
      case <-done: return
      default:
      }

      oldpath := path.Join(pi.installDir, pathToUpdate)
      log.Printf("Updating file %v", oldpath)

      err := pi.backupFile(pathToUpdate)

      if err == nil {
        newpath := path.Join(pi.packageDir, pathToUpdate)
        err = os.Rename(newpath, oldpath)
      }

      if err != nil {
        log.Printf("Updating file %v failed", pathToUpdate)
        log.Println(err)
        errc <- err
        close(done)
      } else {
        go pi.progressReporter.accountUpdate(filesize)
      }
    }()
  }

  go func() {
    errc <- nil
  }()

  wg.Wait()

  if err := <-errc; err != nil {
    return err
  }

  return nil
}

func (pi *PackageInstaller) addFiles(files []*UpdateFileInfo) error {
  log.Printf("Adding %v files", len(files))

  var wg sync.WaitGroup
  errc := make(chan error)
  done := make(chan bool)

  for _, fi := range files {
    wg.Add(1)

    pathToAdd, filesize := fi.Filepath, fi.FileSize

    go func() {
      defer wg.Done()

      select {
      case <-done: return
      default:
      }

      oldpath := path.Join(pi.installDir, pathToAdd)
      ensureDirExists(oldpath)

      newpath := path.Join(pi.packageDir, pathToAdd)
      err := os.Rename(newpath, oldpath)

      if err != nil {
        log.Printf("Adding file %v failed", pathToAdd)
        log.Println(err)
        errc <- err
        close(done)
      } else {
        go pi.progressReporter.accountAdd(filesize)
      }
    }()
  }

  go func() {
    errc <- nil
  }()

  wg.Wait()

  if err := <-errc; err != nil {
    return err
  }

  return nil
}

func purgeFiles(root string, files []*UpdateFileInfo) {
  log.Printf("Purging %v files", len(files))

  var wg sync.WaitGroup

  for _, fi := range files {
    wg.Add(1)

    fileToPurge := fi.Filepath

    go func() {
      defer wg.Done()

      fullpath := path.Join(root, fileToPurge)
      err := os.Remove(fullpath)
      if err != nil {
        log.Println(err)
      }
    }()
  }

  wg.Wait()
}

func ensureDirExists(fullpath string) (err error) {
  dirpath := path.Dir(fullpath)
  err = os.MkdirAll(dirpath, os.ModeDir)
  if err != nil {
    log.Printf("Failed to create directory %v", dirpath)
  }

  return err
}

type ByLength []string

func (s ByLength) Len() int {
    return len(s)
}
func (s ByLength) Swap(i, j int) {
    s[i], s[j] = s[j], s[i]
}
func (s ByLength) Less(i, j int) bool {
    return len(s[i]) > len(s[j])
}

func cleanupEmptyDirs(root string) {
  c := make(chan string)

  go func() {
    var wg sync.WaitGroup
    err := filepath.Walk(root, func(path string, info os.FileInfo, err error) error {
      if err != nil {
        return err
      }

      if info.Mode().IsDir() {
        wg.Add(1)
        go func() {
          c <- path
          wg.Done()
        }()
      }

      return nil
    })

    if err != nil {
      log.Println(err)
    }

    go func() {
      wg.Wait()
      close(c)
    }()
  }()

  dirs := make([]string, 0)
  for path := range c {
    dirs = append(dirs, path)
  }

  removeEmptyDirs(dirs)
}

func removeEmptyDirs(dirs []string) {
  sort.Sort(ByLength(dirs))

  for _, dirpath := range dirs {
    entries, err := ioutil.ReadDir(dirpath)
    if err != nil { continue }

    if len(entries) == 0 {
      log.Printf("Removing empty dir %v", dirpath)

      err = os.Remove(dirpath)
      if err != nil {
        log.Println(err)
      }
    }
  }
}

func (pr *ProgressReporter) accountRemove(progress int64) {
  pr.progressChan <- (progress*RemoveFactor)/100
}

func (pr *ProgressReporter) accountUpdate(progress int64) {
  pr.progressChan <- (progress*UpdateFactor)/100
}

func (pr *ProgressReporter) accountAdd(progress int64) {
  pr.progressChan <- (progress*AddFactor)/100
}

func (pr *ProgressReporter) reportingLoop() {
  for chunk := range pr.progressChan {
    pr.currentProgress += uint64(chunk)

    percent := (pr.currentProgress*100) / pr.grandTotal
    pr.percent = int(percent)

    go func() {
      pr.reportingChan <- true
    }()
  }

  close(pr.reportingChan)
}

func (pr *ProgressReporter) receiveUpdates() {
  for _ = range pr.reportingChan {
    pr.progressHandler.HandlePercentChange(pr.percent)
  }
}

func (pr *ProgressReporter) receiveSystemMessages() {
  for msg := range pr.systemMessageChan {
    pr.progressHandler.HandleSystemMessage(msg)
  }
}

func (pr *ProgressReporter) receiveFinish() {
  <- pr.finished
  pr.progressHandler.HandleFinish()
}

func (pr *ProgressReporter) handleProgress() {
  go pr.receiveSystemMessages()
  go pr.receiveUpdates()
  go pr.receiveFinish()
}

func (ph *LogProgressHandler) HandlePercentChange(percent int) {
  log.Printf("Completed %v%%", percent)
}

func (ph *LogProgressHandler) HandleSystemMessage(msg string) {
  log.Printf("System message: %v", msg)
}

func (ph *LogProgressHandler) HandleFinish() {
  log.Printf("Finished")
}
<|endoftext|>


<commit_before>package api

import (
	"crypto/rand"
	"encoding/base64"
)

func Util_GenerateRandomBytes(n int) ([]byte, error) {
	b := make([]byte, n)
	_, err := rand.Read(b)

	if err != nil {
		return nil, err
	}

	return b, nil
}

func Util_GenerateRandomString(s int) (string, error) {
	b, err := Util_GenerateRandomBytes(s)
	return base64.URLEncoding.EncodeToString(b), err
}<commit_msg>format util.go<commit_after>package api

import (
	"crypto/rand"
	"encoding/base64"
)

func Util_GenerateRandomBytes(n int) ([]byte, error) {
	b := make([]byte, n)
	_, err := rand.Read(b)

	if err != nil {
		return nil, err
	}

	return b, nil
}

func Util_GenerateRandomString(s int) (string, error) {
	b, err := Util_GenerateRandomBytes(s)
	return base64.URLEncoding.EncodeToString(b), err
}
<|endoftext|>


<commit_before>package gopcap

import (
	"bytes"
	"os"
	"testing"
	"time"
)

// Test the overall parsing functionality using the packaged testing .cap file.
func TestParse(t *testing.T) {
	src, err := os.Open("SkypeIRC.cap")
	if err != nil {
		t.Error("Missing pcap file.")
	}

	parsed, err := Parse(src)

	// Check the file header.
	if parsed.MajorVersion != uint16(2) {
		t.Errorf("Incorrectly parsed major version: expected %v, got %v.", 2, parsed.MajorVersion)
	}
	if parsed.MinorVersion != uint16(4) {
		t.Errorf("Incorrectly parsed minor version: expected %v, got %v.", 4, parsed.MinorVersion)
	}
	if parsed.TZCorrection != int32(0) {
		t.Errorf("Got nonzero TZ correction: %v.", parsed.TZCorrection)
	}
	if parsed.SigFigs != uint32(0) {
		t.Errorf("Got nonzero sig figs: %v.", parsed.SigFigs)
	}
	if parsed.MaxLen != uint32(65535) {
		t.Errorf("Incorrectly parsed maximum len: expected %v, got %v.", 65535, parsed.MaxLen)
	}
	if parsed.LinkType != ETHERNET {
		t.Errorf("Incorrect link type: expected %v, got %v.", ETHERNET, parsed.LinkType)
	}
	if len(parsed.Packets) != 2264 {
		t.Errorf("Unexpected number of packets: expected %v, got %v.", 2264, len(parsed.Packets))
	}

	// Check the packet header from the first packet. Including the raw data is a lousy way to test, but
	// at least the packet is small.
	packet := parsed.Packets[0]
	correct_ts := 321259*time.Hour + 31*time.Minute + 6*time.Second + 654*time.Millisecond + 692*time.Microsecond

	if packet.Timestamp != correct_ts {
		t.Errorf("Unexpected TS: expected %v, got %v.", correct_ts, packet.Timestamp)
	}
	if packet.IncludedLen != uint32(96) {
		t.Errorf("Unexpected included length: expected %v, got %v.", 96, packet.IncludedLen)
	}
	if packet.ActualLen != uint32(96) {
		t.Errorf("Unexpected actual length: expected %v, got %v.", 96, packet.ActualLen)
	}

	// This is definitely an ethernet frame. If this fails, we failed the test.
	frame := packet.Data.(*EthernetFrame)
	macSrc := []byte{0x00, 0x04, 0x76, 0x96, 0x7B, 0xDA}
	macDst := []byte{0x00, 0x16, 0xE3, 0x19, 0x27, 0x15}

	if bytes.Compare(frame.MACSource, macSrc) != 0 {
		t.Errorf("Unexpected source MAC: expected %v, got %v.", macSrc, frame.MACSource)
	}
	if bytes.Compare(frame.MACDestination, macDst) != 0 {
		t.Errorf("Unexpected destination MAC: expected %v, got %v.", macDst, frame.MACDestination)
	}
	if len(frame.VLANTag) != 0 {
		t.Errorf("Incorrectly received VLAN tag: %v", frame.VLANTag)
	}
	if frame.Length != 0 {
		t.Errorf("Incorrectly received length: %v", frame.Length)
	}
	if frame.EtherType != EtherType(2048) {
		t.Errorf("Unexpected EtherType: expected %v, got %v", 2048, frame.EtherType)
	}

	// This is definitely an IPv4 packet.
	pkt := frame.LinkData().(*IPv4Packet)
	expectedSrc := []byte{192, 168, 1, 2}
	expectedDst := []byte{212, 204, 214, 114}

	if err != nil {
		t.Errorf("Unexpected error: %v", err)
	}
	if pkt.IHL != uint8(5) {
		t.Errorf("Unexpected IHL: expected %v, got %v", 5, pkt.IHL)
	}
	if pkt.DSCP != uint8(0) {
		t.Errorf("Unexpected DSCP: expected %v, got %v", 0, pkt.DSCP)
	}
	if pkt.ECN != uint8(0) {
		t.Errorf("Unexpected ECN: expected %v, got %v", 0, pkt.ECN)
	}
	if pkt.TotalLength != uint16(82) {
		t.Errorf("Unexpected total length: expected %v, got %v", 82, pkt.TotalLength)
	}
	if pkt.ID != uint16(30445) {
		t.Errorf("Unexpected ID: expected %v, got %v", 30445, pkt.ID)
	}
	if !pkt.DontFragment {
		t.Error("Don't fragment bit unset.")
	}
	if pkt.MoreFragments {
		t.Errorf("More fragments bit set.")
	}
	if pkt.FragmentOffset != uint16(0) {
		t.Errorf("Unexpected fragment offset: expected %v, got %v", 0, pkt.FragmentOffset)
	}
	if pkt.TTL != uint8(64) {
		t.Errorf("Unexpected TTL: expected %v, got %v", 64, pkt.TTL)
	}
	if pkt.Protocol != IPP_TCP {
		t.Errorf("Unexpected protocol: expected %v, got %v", IPP_TCP, pkt.Protocol)
	}
	if pkt.Checksum != uint16(22223) {
		t.Errorf("Unexpected checksum: expected %v, got %v", 22223, pkt.Checksum)
	}
	if bytes.Compare(pkt.SourceAddress, expectedSrc) != 0 {
		t.Errorf("Unexpected source address: expected %v, got %v", expectedSrc, pkt.SourceAddress)
	}
	if bytes.Compare(pkt.DestAddress, expectedDst) != 0 {
		t.Errorf("Unexpected destination address: expected %v, got %v", expectedDst, pkt.DestAddress)
	}
	if len(pkt.Options) != 0 {
		t.Errorf("Shouldn't have any options: got %v", pkt.Options)
	}
}
<commit_msg>Remove superfluous error check.<commit_after>package gopcap

import (
	"bytes"
	"os"
	"testing"
	"time"
)

// Test the overall parsing functionality using the packaged testing .cap file.
func TestParse(t *testing.T) {
	src, err := os.Open("SkypeIRC.cap")
	if err != nil {
		t.Error("Missing pcap file.")
	}

	parsed, err := Parse(src)

	// Check the file header.
	if parsed.MajorVersion != uint16(2) {
		t.Errorf("Incorrectly parsed major version: expected %v, got %v.", 2, parsed.MajorVersion)
	}
	if parsed.MinorVersion != uint16(4) {
		t.Errorf("Incorrectly parsed minor version: expected %v, got %v.", 4, parsed.MinorVersion)
	}
	if parsed.TZCorrection != int32(0) {
		t.Errorf("Got nonzero TZ correction: %v.", parsed.TZCorrection)
	}
	if parsed.SigFigs != uint32(0) {
		t.Errorf("Got nonzero sig figs: %v.", parsed.SigFigs)
	}
	if parsed.MaxLen != uint32(65535) {
		t.Errorf("Incorrectly parsed maximum len: expected %v, got %v.", 65535, parsed.MaxLen)
	}
	if parsed.LinkType != ETHERNET {
		t.Errorf("Incorrect link type: expected %v, got %v.", ETHERNET, parsed.LinkType)
	}
	if len(parsed.Packets) != 2264 {
		t.Errorf("Unexpected number of packets: expected %v, got %v.", 2264, len(parsed.Packets))
	}

	// Check the packet header from the first packet. Including the raw data is a lousy way to test, but
	// at least the packet is small.
	packet := parsed.Packets[0]
	correct_ts := 321259*time.Hour + 31*time.Minute + 6*time.Second + 654*time.Millisecond + 692*time.Microsecond

	if packet.Timestamp != correct_ts {
		t.Errorf("Unexpected TS: expected %v, got %v.", correct_ts, packet.Timestamp)
	}
	if packet.IncludedLen != uint32(96) {
		t.Errorf("Unexpected included length: expected %v, got %v.", 96, packet.IncludedLen)
	}
	if packet.ActualLen != uint32(96) {
		t.Errorf("Unexpected actual length: expected %v, got %v.", 96, packet.ActualLen)
	}

	// This is definitely an ethernet frame. If this fails, we failed the test.
	frame := packet.Data.(*EthernetFrame)
	macSrc := []byte{0x00, 0x04, 0x76, 0x96, 0x7B, 0xDA}
	macDst := []byte{0x00, 0x16, 0xE3, 0x19, 0x27, 0x15}

	if bytes.Compare(frame.MACSource, macSrc) != 0 {
		t.Errorf("Unexpected source MAC: expected %v, got %v.", macSrc, frame.MACSource)
	}
	if bytes.Compare(frame.MACDestination, macDst) != 0 {
		t.Errorf("Unexpected destination MAC: expected %v, got %v.", macDst, frame.MACDestination)
	}
	if len(frame.VLANTag) != 0 {
		t.Errorf("Incorrectly received VLAN tag: %v", frame.VLANTag)
	}
	if frame.Length != 0 {
		t.Errorf("Incorrectly received length: %v", frame.Length)
	}
	if frame.EtherType != EtherType(2048) {
		t.Errorf("Unexpected EtherType: expected %v, got %v", 2048, frame.EtherType)
	}

	// This is definitely an IPv4 packet.
	pkt := frame.LinkData().(*IPv4Packet)
	expectedSrc := []byte{192, 168, 1, 2}
	expectedDst := []byte{212, 204, 214, 114}

	if pkt.IHL != uint8(5) {
		t.Errorf("Unexpected IHL: expected %v, got %v", 5, pkt.IHL)
	}
	if pkt.DSCP != uint8(0) {
		t.Errorf("Unexpected DSCP: expected %v, got %v", 0, pkt.DSCP)
	}
	if pkt.ECN != uint8(0) {
		t.Errorf("Unexpected ECN: expected %v, got %v", 0, pkt.ECN)
	}
	if pkt.TotalLength != uint16(82) {
		t.Errorf("Unexpected total length: expected %v, got %v", 82, pkt.TotalLength)
	}
	if pkt.ID != uint16(30445) {
		t.Errorf("Unexpected ID: expected %v, got %v", 30445, pkt.ID)
	}
	if !pkt.DontFragment {
		t.Error("Don't fragment bit unset.")
	}
	if pkt.MoreFragments {
		t.Errorf("More fragments bit set.")
	}
	if pkt.FragmentOffset != uint16(0) {
		t.Errorf("Unexpected fragment offset: expected %v, got %v", 0, pkt.FragmentOffset)
	}
	if pkt.TTL != uint8(64) {
		t.Errorf("Unexpected TTL: expected %v, got %v", 64, pkt.TTL)
	}
	if pkt.Protocol != IPP_TCP {
		t.Errorf("Unexpected protocol: expected %v, got %v", IPP_TCP, pkt.Protocol)
	}
	if pkt.Checksum != uint16(22223) {
		t.Errorf("Unexpected checksum: expected %v, got %v", 22223, pkt.Checksum)
	}
	if bytes.Compare(pkt.SourceAddress, expectedSrc) != 0 {
		t.Errorf("Unexpected source address: expected %v, got %v", expectedSrc, pkt.SourceAddress)
	}
	if bytes.Compare(pkt.DestAddress, expectedDst) != 0 {
		t.Errorf("Unexpected destination address: expected %v, got %v", expectedDst, pkt.DestAddress)
	}
	if len(pkt.Options) != 0 {
		t.Errorf("Shouldn't have any options: got %v", pkt.Options)
	}
}
<|endoftext|>


<commit_before>package gorm

import (
	"strconv"
	"time"
	"log"

	db "github.com/theduke/go-dukedb"
)

/**
 * Implement migration related interfaces.
 */
func (b Backend) GetMigrationHandler() *db.MigrationHandler {
	return b.MigrationHandler
}

func (b Backend) MigrationsSetup() db.DbError {
	count := -1
	b.Db.Model(MigrationAttempt{}).Count(&count)
	initalRun := count == -1

	if initalRun {
		log.Println("MIGRATE: Building migration tables.")
		tx := b.Db.Begin()

		if err := tx.CreateTable(MigrationAttempt{}).Error; err != nil {
			return db.Error{
				Code: "migration_setup_failed",
				Message: "Could not create migrations table: " + err.Error(),
				Data: err,
			}

			tx.Rollback()
			return db.Error{
				Code: "migration_setup_failed",
				Message: err.Error(),
			}
		}

		migration := MigrationAttempt{
			Version:    0,
			StartedAt:  time.Now(),
			FinishedAt: time.Now(),
			Complete:  true,
		}
		if err := tx.Create(&migration).Error; err != nil {
			return db.Error{
				Code: "migration_setup_failed",
				Message: "Could not create migrations table: " + err.Error(),
				Data: err,
			}
		}

		tx.Commit()
		log.Println("MIGRATE: Migrations table created.")
	}

	return nil
}

func (b Backend) IsMigrationLocked() (bool, db.DbError) {
	var lastAttempt MigrationAttempt
	if err := b.Db.Last(&lastAttempt).Error; err != nil {
		return true, db.Error{
			Code: "db_error",
			Message: err.Error(),
		}
	}

	if lastAttempt.ID != 0 && lastAttempt.FinishedAt.IsZero() {
		// Last attempt was aborted. DB is locked.
		return true, nil
	}
	return false, nil
}

func (b Backend) DetermineMigrationVersion() (int, db.DbError) {
	var lastAttempt MigrationAttempt
	if err := b.Db.Where("complete = ?", true).Last(&lastAttempt).Error; err != nil {
		return -1, db.Error{
			Code: "db_error",
			Message: err.Error(),
		}
	}

	return lastAttempt.Version, nil
}

func (b Backend) NewMigrationAttempt() db.MigrationAttempt {
	return &MigrationAttempt{}
}

type MigrationAttempt struct {
	ID uint64
	Version int
	StartedAt time.Time
	FinishedAt time.Time
	Complete bool
}

func (m MigrationAttempt) GetCollection() string {
	return "migration_attempts"
}

func(a *MigrationAttempt) GetID() string {
	return strconv.FormatUint(a.ID, 10)
}

func(a *MigrationAttempt) SetID(x string) error {
	id, err := strconv.ParseUint(x, 10, 64)
	if err != nil {
		return err
	}
	a.ID = id
	return nil
}

func(a *MigrationAttempt) GetVersion() int {
	return a.Version
}

func(a *MigrationAttempt) SetVersion(x int) {
	a.Version = x
}

func(a *MigrationAttempt) GetStartedAt() time.Time {
	return a.StartedAt
}

func(a *MigrationAttempt) SetStartedAt(x time.Time) {
	a.StartedAt = x
}

func(a *MigrationAttempt) GetFinishedAt() time.Time {
	return a.FinishedAt
}

func(a *MigrationAttempt) SetFinishedAt(x time.Time) {
	a.FinishedAt = x
}

func(a *MigrationAttempt) GetComplete() bool {
	return a.Complete
}

func(a *MigrationAttempt) SetComplete(x bool) {
	a.Complete = x
}

/*
	// Check if the migration attempts table exists.
	// Otherwise, create it and run the additional migration.

	// Determine if the database is locked.

	// Determine current version of the database.

	if curVersion < targetVersion {
		for nextVersion := curVersion + 1; nextVersion <= targetVersion; nextVersion++ {
			migration := m.Get(nextVersion)
			if err := migration.Run(db); err != nil {
				// Migration failed! Abort.
				return err
			}
		}
	} else {
		log.Println("DB is already at newest schema version: " + strconv.Itoa(targetVersion))
	}
*/<commit_msg>now using basemigrationattempt in gorm backend<commit_after>package gorm

import (
	"time"
	"log"

	db "github.com/theduke/go-dukedb"
)

/**
 * Implement migration related interfaces.
 */
func (b Backend) GetMigrationHandler() *db.MigrationHandler {
	return b.MigrationHandler
}

func (b Backend) MigrationsSetup() db.DbError {
	count := -1
	b.Db.Model(MigrationAttempt{}).Count(&count)
	initalRun := count == -1

	if initalRun {
		log.Println("MIGRATE: Building migration tables.")
		tx := b.Db.Begin()

		if err := tx.CreateTable(MigrationAttempt{}).Error; err != nil {
			return db.Error{
				Code: "migration_setup_failed",
				Message: "Could not create migrations table: " + err.Error(),
				Data: err,
			}

			tx.Rollback()
			return db.Error{
				Code: "migration_setup_failed",
				Message: err.Error(),
			}
		}

		migration := MigrationAttempt{}
		migration.Version = 0
		migration.StartedAt = time.Now()
		migration.FinishedAt = time.Now()
		migration.Complete = true

		if err := tx.Create(&migration).Error; err != nil {
			return db.Error{
				Code: "migration_setup_failed",
				Message: "Could not create migrations table: " + err.Error(),
				Data: err,
			}
		}

		tx.Commit()
		log.Println("MIGRATE: Migrations table created.")
	}

	return nil
}

func (b Backend) IsMigrationLocked() (bool, db.DbError) {
	var lastAttempt MigrationAttempt
	if err := b.Db.Last(&lastAttempt).Error; err != nil {
		return true, db.Error{
			Code: "db_error",
			Message: err.Error(),
		}
	}

	if lastAttempt.ID != 0 && lastAttempt.FinishedAt.IsZero() {
		// Last attempt was aborted. DB is locked.
		return true, nil
	}
	return false, nil
}

func (b Backend) DetermineMigrationVersion() (int, db.DbError) {
	var lastAttempt MigrationAttempt
	if err := b.Db.Where("complete = ?", true).Last(&lastAttempt).Error; err != nil {
		return -1, db.Error{
			Code: "db_error",
			Message: err.Error(),
		}
	}

	return lastAttempt.Version, nil
}

type MigrationAttempt struct {
	db.BaseMigrationAttemptIntID
}

func (b Backend) NewMigrationAttempt() db.MigrationAttempt {
	return &MigrationAttempt{}
}<|endoftext|>


<commit_before>// HttpKnock
// dRbiG
// See LICENSE.txt

package main

import (
	"flag"
	"fmt"
	"log"
	"net/http"
	"os"
	"os/exec"
	"os/signal"
	"strings"
	"sync"
	"syscall"
	"time"
)

const (
	//	OPEN_FW_CMD_FMT  = `sudo ipfw table 1 add %s`
	//	CLOSE_FW_CMD_FMT = `sudo ipfw table 1 delete %s`
	OPEN_FW_CMD_FMT  = `echo open %s`
	CLOSE_FW_CMD_FMT = `echo close %s`
	KEEP_DURATION    = time.Duration(5) * time.Second
	HOST             = `0.0.0.0`
	PORT             = 9996
	PASSWORD_VAR     = `HK_PASSWORD`
	PASSWORD_KEY     = `key`
)

const (
	VERSION  = `0.1`
	HELP_FMT = `Usage: %s (options)
httpknock v%s, see LICENSE.txt

Set password using %s env variable.

`
)

type timerMap struct {
	mu sync.Mutex
	ts map[string]*time.Timer
}

var (
	flagKeepDuration time.Duration
	flagHost         string
	flagPort         int
)

var (
	password string
	timers   timerMap
)

func init() {
	flag.Usage = func() {
		fmt.Fprintf(os.Stderr, HELP_FMT, os.Args[0], VERSION, PASSWORD_VAR)
		flag.PrintDefaults()
	}

	flag.DurationVar(&flagKeepDuration, "kd", KEEP_DURATION, "Keep open for given duration.")
	flag.StringVar(&flagHost, "h", HOST, "Host to bind to.")
	flag.IntVar(&flagPort, "p", PORT, "Port to bind to.")

	timers.ts = make(map[string]*time.Timer, 256)
}

func main() {
	var ok bool

	if password, ok = get_password(); !ok {
		fmt.Fprintf(os.Stderr, "Variable %s not set. Can't run without password.\n", PASSWORD_VAR)
		os.Exit(1)
	}

	flag.Parse()
	go runHTTPServer()
	sigwait()

	log.Println("HttpKnock stopped.")
}

func handleOpen(w http.ResponseWriter, req *http.Request) {
	if !auth_request(w, req) {
		return
	}

	ip := get_ip(req.RemoteAddr)
	if !run_fw_cmd(OPEN_FW_CMD_FMT, ip) {
		fmt.Fprintln(w, "FAILED")
		return
	}

	duration := flagKeepDuration
	if val := req.FormValue("for"); val != "" {
		user_duration, err := time.ParseDuration(val)
		if err != nil {
			log.Printf("Failed to parse duration: %s", err)
			fmt.Fprintln(w, "Failed to parse duration, using default.")
		} else {
			duration = user_duration
		}
	}

	timers.mu.Lock()
	if t, ok := timers.ts[ip]; ok {
		if !t.Stop() {
			<-t.C
		}
		t.Reset(duration)

		until := time.Now().Add(duration)
		log.Printf("Client %s reset timer for %s to: %s", req.RemoteAddr, ip, until)
		fmt.Fprintf(w, "Reset timer for %s to: %s\n", ip, until)
		fmt.Fprintln(w, "OK")
		timers.mu.Unlock()
		return
	}

	timers.ts[ip] = time.AfterFunc(duration, func() {
		log.Printf("Closing FW for %s after %s timeout...", ip, duration)
		run_fw_cmd(CLOSE_FW_CMD_FMT, ip)
		timers.mu.Lock()
		delete(timers.ts, ip)
		timers.mu.Unlock()
	})
	timers.mu.Unlock()

	until := time.Now().Add(duration)
	log.Printf("Added %s until %s\n", ip, until)

	fmt.Fprintf(w, "Added %s until %s.\n", ip, until)
	fmt.Fprintln(w, "OK")
}

func handleClose(w http.ResponseWriter, req *http.Request) {
	if !auth_request(w, req) {
		return
	}

	var ip string
	if val := req.FormValue("ip"); val != "" {
		ip = val
	} else {
		ip = get_ip(req.RemoteAddr)
	}

	timers.mu.Lock()
	if _, ok := timers.ts[ip]; ok {
		if !run_fw_cmd(CLOSE_FW_CMD_FMT, ip) {
			fmt.Fprintln(w, "FAILED")
			timers.mu.Unlock()
			return
		}

		timers.ts[ip].Stop()
		delete(timers.ts, ip)

		log.Printf("Client %s killed timer for: %s\n", req.RemoteAddr, ip)
		fmt.Fprintln(w, "OK")
	} else {
		log.Printf("Client %s tried to unblock non-blocked ip: %s\n", req.RemoteAddr, ip)
		fmt.Fprintf(w, "IP %s is not open.\n", ip)
		fmt.Fprintln(w, "FAILED")
	}

	timers.mu.Unlock()
}

func run_fw_cmd(cmd_fmt string, addr string) bool {
	cmd_str := fmt.Sprintf(cmd_fmt, addr)
	cmd_args := strings.Split(cmd_str, " ")

	if output, err := exec.Command(cmd_args[0], cmd_args[1:]...).Output(); err != nil {
		log.Printf("Command '%s' failed: %s\nOutput: %s\n", cmd_str, err, string(output))
		return false
	}

	log.Printf("Command '%s' succeeded", cmd_str)
	return true
}

func runHTTPServer() {
	addr := fmt.Sprintf("%s:%d", flagHost, flagPort)
	http.HandleFunc("/open", handleOpen)
	http.HandleFunc("/close", handleClose)
	log.Println("Starting HTTP server at", addr)
	log.Fatalln(http.ListenAndServe(addr, nil))
}

func auth_request(w http.ResponseWriter, req *http.Request) bool {
	if val := req.FormValue(PASSWORD_KEY); val == password {
		log.Printf("Authorized access to %s form %s\n", req.RequestURI, req.RemoteAddr)
		return true
	}

	log.Printf("Unauthorized access to %s from %s\n", req.RequestURI, req.RemoteAddr)
	http.NotFound(w, req)
	return false
}

func get_password() (string, bool) {
	for _, e := range os.Environ() {
		kv := strings.Split(e, "=")
		if kv[0] == PASSWORD_VAR {
			return kv[1], true
		}
	}

	return "", false
}

func get_ip(addr string) string {
	return strings.Split(addr, ":")[0]
}

func sigwait() {
	sig := make(chan os.Signal)
	signal.Notify(sig, syscall.SIGINT, syscall.SIGTERM, syscall.SIGKILL)

	stop_sig := <-sig
	log.Printf("Received signal: %s\n", stop_sig)
}
<commit_msg>httpknock.go: Add info stuff<commit_after>// HttpKnock
// dRbiG
// See LICENSE.txt

package main

import (
	"flag"
	"fmt"
	"log"
	"net/http"
	"os"
	"os/exec"
	"os/signal"
	"strings"
	"sync"
	"syscall"
	"time"
)

const (
	//	OPEN_FW_CMD_FMT  = `sudo ipfw table 1 add %s`
	//	CLOSE_FW_CMD_FMT = `sudo ipfw table 1 delete %s`
	OPEN_FW_CMD_FMT  = `echo open %s`
	CLOSE_FW_CMD_FMT = `echo close %s`
	KEEP_DURATION    = time.Duration(5) * time.Second
	HOST             = `0.0.0.0`
	PORT             = 9996
	PASSWORD_VAR     = `HK_PASSWORD`
	PASSWORD_KEY     = `key`
)

const (
	VERSION  = `0.1`
	HELP_FMT = `Usage: %s (options)
httpknock v%s, see LICENSE.txt

Set password using %s env variable.

`
)

type timerEntry struct {
	t *time.Timer
	u time.Time
}

type timerMap struct {
	mu sync.Mutex
	ts map[string]timerEntry
}

var (
	flagKeepDuration time.Duration
	flagHost         string
	flagPort         int
)

var (
	password string
	timers   timerMap
)

func init() {
	flag.Usage = func() {
		fmt.Fprintf(os.Stderr, HELP_FMT, os.Args[0], VERSION, PASSWORD_VAR)
		flag.PrintDefaults()
	}

	flag.DurationVar(&flagKeepDuration, "kd", KEEP_DURATION, "Keep open for given duration.")
	flag.StringVar(&flagHost, "h", HOST, "Host to bind to.")
	flag.IntVar(&flagPort, "p", PORT, "Port to bind to.")

	timers.ts = make(map[string]timerEntry, 256)
}

func main() {
	var ok bool

	if password, ok = get_password(); !ok {
		fmt.Fprintf(os.Stderr, "Variable %s not set. Can't run without password.\n", PASSWORD_VAR)
		os.Exit(1)
	}

	flag.Parse()
	go runHTTPServer()
	sigwait()

	log.Println("HttpKnock stopped.")
}

func handleOpen(w http.ResponseWriter, req *http.Request) {
	if !auth_request(w, req) {
		return
	}

	ip := get_ip(req.RemoteAddr)
	if !run_fw_cmd(OPEN_FW_CMD_FMT, ip) {
		fmt.Fprintln(w, "FAILED")
		return
	}

	duration := flagKeepDuration
	if val := req.FormValue("for"); val != "" {
		user_duration, err := time.ParseDuration(val)
		if err != nil {
			log.Printf("Failed to parse duration: %s", err)
			fmt.Fprintln(w, "Failed to parse duration, using default.")
		} else {
			duration = user_duration
		}
	}
	until := time.Now().Add(duration)

	timers.mu.Lock()
	if t, ok := timers.ts[ip]; ok {
		if !t.t.Stop() {
			<-t.t.C
		}
		t.t.Reset(duration)

		until := time.Now().Add(duration)
		log.Printf("Client %s reset timer for %s to: %s", req.RemoteAddr, ip, until)
		fmt.Fprintf(w, "Reset timer for %s to: %s\n", ip, until)
		fmt.Fprintln(w, "OK")
		timers.mu.Unlock()
		return
	}

	timers.ts[ip] = timerEntry{
		u: until,
		t: time.AfterFunc(duration, func() {
			log.Printf("Closing FW for %s after %s timeout...", ip, duration)
			run_fw_cmd(CLOSE_FW_CMD_FMT, ip)
			timers.mu.Lock()
			delete(timers.ts, ip)
			timers.mu.Unlock()
		}),
	}
	timers.mu.Unlock()

	log.Printf("Added %s until %s\n", ip, until)

	fmt.Fprintf(w, "Added %s until %s.\n", ip, until)
	fmt.Fprintln(w, "OK")
}

func handleClose(w http.ResponseWriter, req *http.Request) {
	if !auth_request(w, req) {
		return
	}

	var ip string
	if val := req.FormValue("ip"); val != "" {
		ip = val
	} else {
		ip = get_ip(req.RemoteAddr)
	}

	timers.mu.Lock()
	if _, ok := timers.ts[ip]; ok {
		if !run_fw_cmd(CLOSE_FW_CMD_FMT, ip) {
			fmt.Fprintln(w, "FAILED")
			timers.mu.Unlock()
			return
		}

		timers.ts[ip].t.Stop()
		delete(timers.ts, ip)

		log.Printf("Client %s killed timer for: %s\n", req.RemoteAddr, ip)
		fmt.Fprintln(w, "OK")
	} else {
		log.Printf("Client %s tried to unblock non-blocked ip: %s\n", req.RemoteAddr, ip)
		fmt.Fprintf(w, "IP %s is not open.\n", ip)
		fmt.Fprintln(w, "FAILED")
	}

	timers.mu.Unlock()
}

func handleInfo(w http.ResponseWriter, req *http.Request) {
	if !auth_request(w, req) {
		return
	}

	fmt.Fprintln(w, "IP      \t\tExpires")
	timers.mu.Lock()
	for ip, t := range timers.ts {
		fmt.Fprintf(w, "%s\t\t%s\n", ip, t.u)
	}
	timers.mu.Unlock()
	fmt.Fprintln(w, "OK")
}

func run_fw_cmd(cmd_fmt string, addr string) bool {
	cmd_str := fmt.Sprintf(cmd_fmt, addr)
	cmd_args := strings.Split(cmd_str, " ")

	if output, err := exec.Command(cmd_args[0], cmd_args[1:]...).Output(); err != nil {
		log.Printf("Command '%s' failed: %s\nOutput: %s\n", cmd_str, err, string(output))
		return false
	}

	log.Printf("Command '%s' succeeded", cmd_str)
	return true
}

func runHTTPServer() {
	addr := fmt.Sprintf("%s:%d", flagHost, flagPort)
	http.HandleFunc("/open", handleOpen)
	http.HandleFunc("/close", handleClose)
	http.HandleFunc("/info", handleInfo)
	log.Println("Starting HTTP server at", addr)
	log.Fatalln(http.ListenAndServe(addr, nil))
}

func auth_request(w http.ResponseWriter, req *http.Request) bool {
	if val := req.FormValue(PASSWORD_KEY); val == password {
		log.Printf("Authorized access to %s form %s\n", req.RequestURI, req.RemoteAddr)
		return true
	}

	log.Printf("Unauthorized access to %s from %s\n", req.RequestURI, req.RemoteAddr)
	http.NotFound(w, req)
	return false
}

func get_password() (string, bool) {
	for _, e := range os.Environ() {
		kv := strings.Split(e, "=")
		if kv[0] == PASSWORD_VAR {
			return kv[1], true
		}
	}

	return "", false
}

func get_ip(addr string) string {
	return strings.Split(addr, ":")[0]
}

func sigwait() {
	sig := make(chan os.Signal)
	signal.Notify(sig, syscall.SIGINT, syscall.SIGTERM, syscall.SIGKILL)

	stop_sig := <-sig
	log.Printf("Received signal: %s\n", stop_sig)
}
<|endoftext|>


<commit_before>/*
Copyright 2014 Google Inc. All rights reserved.

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

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

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

package api

import (
	"strings"

	errs "github.com/GoogleCloudPlatform/kubernetes/pkg/api/errors"
	"github.com/GoogleCloudPlatform/kubernetes/pkg/labels"
	"github.com/GoogleCloudPlatform/kubernetes/pkg/util"
)

func validateVolumes(volumes []Volume) (util.StringSet, errs.ErrorList) {
	allErrs := errs.ErrorList{}

	allNames := util.StringSet{}
	for i := range volumes {
		vol := &volumes[i] // so we can set default values
		el := errs.ErrorList{}
		// TODO(thockin) enforce that a source is set once we deprecate the implied form.
		if vol.Source != nil {
			el = validateSource(vol.Source).Prefix("source")
		}
		if len(vol.Name) == 0 {
			el = append(el, errs.NewRequired("name", vol.Name))
		} else if !util.IsDNSLabel(vol.Name) {
			el = append(el, errs.NewInvalid("name", vol.Name))
		} else if allNames.Has(vol.Name) {
			el = append(el, errs.NewDuplicate("name", vol.Name))
		}
		if len(el) == 0 {
			allNames.Insert(vol.Name)
		} else {
			allErrs = append(allErrs, el.PrefixIndex(i)...)
		}
	}
	return allNames, allErrs
}

func validateSource(source *VolumeSource) errs.ErrorList {
	numVolumes := 0
	allErrs := errs.ErrorList{}
	if source.HostDirectory != nil {
		numVolumes++
		allErrs = append(allErrs, validateHostDir(source.HostDirectory).Prefix("hostDirectory")...)
	}
	if source.EmptyDirectory != nil {
		numVolumes++
		//EmptyDirs have nothing to validate
	}
	if numVolumes != 1 {
		allErrs = append(allErrs, errs.NewInvalid("", source))
	}
	return allErrs
}

func validateHostDir(hostDir *HostDirectory) errs.ErrorList {
	allErrs := errs.ErrorList{}
	if hostDir.Path == "" {
		allErrs = append(allErrs, errs.NewNotFound("path", hostDir.Path))
	}
	return allErrs
}

var supportedPortProtocols = util.NewStringSet("TCP", "UDP")

func validatePorts(ports []Port) errs.ErrorList {
	allErrs := errs.ErrorList{}

	allNames := util.StringSet{}
	for i := range ports {
		pErrs := errs.ErrorList{}
		port := &ports[i] // so we can set default values
		if len(port.Name) > 0 {
			if len(port.Name) > 63 || !util.IsDNSLabel(port.Name) {
				pErrs = append(pErrs, errs.NewInvalid("name", port.Name))
			} else if allNames.Has(port.Name) {
				pErrs = append(pErrs, errs.NewDuplicate("name", port.Name))
			} else {
				allNames.Insert(port.Name)
			}
		}
		if port.ContainerPort == 0 {
			pErrs = append(pErrs, errs.NewRequired("containerPort", port.ContainerPort))
		} else if !util.IsValidPortNum(port.ContainerPort) {
			pErrs = append(pErrs, errs.NewInvalid("containerPort", port.ContainerPort))
		}
		if port.HostPort != 0 && !util.IsValidPortNum(port.HostPort) {
			pErrs = append(pErrs, errs.NewInvalid("hostPort", port.HostPort))
		}
		if len(port.Protocol) == 0 {
			port.Protocol = "TCP"
		} else if !supportedPortProtocols.Has(strings.ToUpper(port.Protocol)) {
			pErrs = append(pErrs, errs.NewNotSupported("protocol", port.Protocol))
		}
		allErrs = append(allErrs, pErrs.PrefixIndex(i)...)
	}
	return allErrs
}

func validateEnv(vars []EnvVar) errs.ErrorList {
	allErrs := errs.ErrorList{}

	for i := range vars {
		vErrs := errs.ErrorList{}
		ev := &vars[i] // so we can set default values
		if len(ev.Name) == 0 {
			vErrs = append(vErrs, errs.NewRequired("name", ev.Name))
		}
		if !util.IsCIdentifier(ev.Name) {
			vErrs = append(vErrs, errs.NewInvalid("name", ev.Name))
		}
		allErrs = append(allErrs, vErrs.PrefixIndex(i)...)
	}
	return allErrs
}

func validateVolumeMounts(mounts []VolumeMount, volumes util.StringSet) errs.ErrorList {
	allErrs := errs.ErrorList{}

	for i := range mounts {
		mErrs := errs.ErrorList{}
		mnt := &mounts[i] // so we can set default values
		if len(mnt.Name) == 0 {
			mErrs = append(mErrs, errs.NewRequired("name", mnt.Name))
		} else if !volumes.Has(mnt.Name) {
			mErrs = append(mErrs, errs.NewNotFound("name", mnt.Name))
		}
		if len(mnt.MountPath) == 0 {
			mErrs = append(mErrs, errs.NewRequired("mountPath", mnt.MountPath))
		}
		allErrs = append(allErrs, mErrs.PrefixIndex(i)...)
	}
	return allErrs
}

// AccumulateUniquePorts runs an extraction function on each Port of each Container,
// accumulating the results and returning an error if any ports conflict.
func AccumulateUniquePorts(containers []Container, accumulator map[int]bool, extract func(*Port) int) errs.ErrorList {
	allErrs := errs.ErrorList{}

	for ci := range containers {
		cErrs := errs.ErrorList{}
		ctr := &containers[ci]
		for pi := range ctr.Ports {
			port := extract(&ctr.Ports[pi])
			if port == 0 {
				continue
			}
			if accumulator[port] {
				cErrs = append(cErrs, errs.NewDuplicate("Port", port))
			} else {
				accumulator[port] = true
			}
		}
		allErrs = append(allErrs, cErrs.PrefixIndex(ci)...)
	}
	return allErrs
}

// Checks for colliding Port.HostPort values across a slice of containers.
func checkHostPortConflicts(containers []Container) errs.ErrorList {
	allPorts := map[int]bool{}
	return AccumulateUniquePorts(containers, allPorts, func(p *Port) int { return p.HostPort })
}

func validateContainers(containers []Container, volumes util.StringSet) errs.ErrorList {
	allErrs := errs.ErrorList{}

	allNames := util.StringSet{}
	for i := range containers {
		cErrs := errs.ErrorList{}
		ctr := &containers[i] // so we can set default values
		if len(ctr.Name) == 0 {
			cErrs = append(cErrs, errs.NewRequired("name", ctr.Name))
		} else if !util.IsDNSLabel(ctr.Name) {
			cErrs = append(cErrs, errs.NewInvalid("name", ctr.Name))
		} else if allNames.Has(ctr.Name) {
			cErrs = append(cErrs, errs.NewDuplicate("name", ctr.Name))
		} else {
			allNames.Insert(ctr.Name)
		}
		if len(ctr.Image) == 0 {
			cErrs = append(cErrs, errs.NewInvalid("image", ctr.Name))
		}
		cErrs = append(cErrs, validatePorts(ctr.Ports).Prefix("ports")...)
		cErrs = append(cErrs, validateEnv(ctr.Env).Prefix("env")...)
		cErrs = append(cErrs, validateVolumeMounts(ctr.VolumeMounts, volumes).Prefix("volumeMounts")...)
		allErrs = append(allErrs, cErrs.PrefixIndex(i)...)
	}
	// Check for colliding ports across all containers.
	// TODO(thockin): This really is dependent on the network config of the host (IP per pod?)
	// and the config of the new manifest.  But we have not specced that out yet, so we'll just
	// make some assumptions for now.  As of now, pods share a network namespace, which means that
	// every Port.HostPort across the whole pod must be unique.
	allErrs = append(allErrs, checkHostPortConflicts(containers)...)

	return allErrs
}

var supportedManifestVersions = util.NewStringSet("v1beta1", "v1beta2")

// ValidateManifest tests that the specified ContainerManifest has valid data.
// This includes checking formatting and uniqueness.  It also canonicalizes the
// structure by setting default values and implementing any backwards-compatibility
// tricks.
func ValidateManifest(manifest *ContainerManifest) errs.ErrorList {
	allErrs := errs.ErrorList{}

	if len(manifest.Version) == 0 {
		allErrs = append(allErrs, errs.NewRequired("version", manifest.Version))
	} else if !supportedManifestVersions.Has(strings.ToLower(manifest.Version)) {
		allErrs = append(allErrs, errs.NewNotSupported("version", manifest.Version))
	}
	allVolumes, errs := validateVolumes(manifest.Volumes)
	allErrs = append(allErrs, errs.Prefix("volumes")...)
	allErrs = append(allErrs, validateContainers(manifest.Containers, allVolumes).Prefix("containers")...)
	return allErrs
}

func ValidatePodState(podState *PodState) errs.ErrorList {
	allErrs := errs.ErrorList(ValidateManifest(&podState.Manifest)).Prefix("manifest")
	if podState.RestartPolicy.Type == "" {
		podState.RestartPolicy.Type = RestartAlways
	} else if podState.RestartPolicy.Type != RestartAlways &&
		podState.RestartPolicy.Type != RestartOnFailure &&
		podState.RestartPolicy.Type != RestartNever {
		allErrs = append(allErrs, errs.NewNotSupported("restartPolicy.type", podState.RestartPolicy.Type))
	}

	return allErrs
}

// Pod tests if required fields in the pod are set.
func ValidatePod(pod *Pod) errs.ErrorList {
	allErrs := errs.ErrorList{}
	if len(pod.ID) == 0 {
		allErrs = append(allErrs, errs.NewRequired("id", pod.ID))
	}
	allErrs = append(allErrs, ValidatePodState(&pod.DesiredState).Prefix("desiredState")...)
	return allErrs
}

// ValidateService tests if required fields in the service are set.
func ValidateService(service *Service) errs.ErrorList {
	allErrs := errs.ErrorList{}
	if len(service.ID) == 0 {
		allErrs = append(allErrs, errs.NewRequired("id", service.ID))
	} else if !util.IsDNS952Label(service.ID) {
		allErrs = append(allErrs, errs.NewInvalid("id", service.ID))
	}
	if !util.IsValidPortNum(service.Port) {
		allErrs = append(allErrs, errs.NewInvalid("Service.Port", service.Port))
	}
	if labels.Set(service.Selector).AsSelector().Empty() {
		allErrs = append(allErrs, errs.NewRequired("selector", service.Selector))
	}
	return allErrs
}

// ValidateReplicationController tests if required fields in the replication controller are set.
func ValidateReplicationController(controller *ReplicationController) errs.ErrorList {
	allErrs := errs.ErrorList{}
	if len(controller.ID) == 0 {
		allErrs = append(allErrs, errs.NewRequired("id", controller.ID))
	}
	if labels.Set(controller.DesiredState.ReplicaSelector).AsSelector().Empty() {
		allErrs = append(allErrs, errs.NewRequired("desiredState.replicaSelector", controller.DesiredState.ReplicaSelector))
	}
	selector := labels.Set(controller.DesiredState.ReplicaSelector).AsSelector()
	labels := labels.Set(controller.DesiredState.PodTemplate.Labels)
	if !selector.Matches(labels) {
		allErrs = append(allErrs, errs.NewInvalid("desiredState.podTemplate.labels", controller.DesiredState.PodTemplate))
	}
	if controller.DesiredState.Replicas < 0 {
		allErrs = append(allErrs, errs.NewInvalid("desiredState.replicas", controller.DesiredState.Replicas))
	}
	allErrs = append(allErrs, ValidateManifest(&controller.DesiredState.PodTemplate.DesiredState.Manifest).Prefix("desiredState.podTemplate.desiredState.manifest")...)
	return allErrs
}
<commit_msg>Incorrect validation error for container image<commit_after>/*
Copyright 2014 Google Inc. All rights reserved.

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

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

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

package api

import (
	"strings"

	errs "github.com/GoogleCloudPlatform/kubernetes/pkg/api/errors"
	"github.com/GoogleCloudPlatform/kubernetes/pkg/labels"
	"github.com/GoogleCloudPlatform/kubernetes/pkg/util"
)

func validateVolumes(volumes []Volume) (util.StringSet, errs.ErrorList) {
	allErrs := errs.ErrorList{}

	allNames := util.StringSet{}
	for i := range volumes {
		vol := &volumes[i] // so we can set default values
		el := errs.ErrorList{}
		// TODO(thockin) enforce that a source is set once we deprecate the implied form.
		if vol.Source != nil {
			el = validateSource(vol.Source).Prefix("source")
		}
		if len(vol.Name) == 0 {
			el = append(el, errs.NewRequired("name", vol.Name))
		} else if !util.IsDNSLabel(vol.Name) {
			el = append(el, errs.NewInvalid("name", vol.Name))
		} else if allNames.Has(vol.Name) {
			el = append(el, errs.NewDuplicate("name", vol.Name))
		}
		if len(el) == 0 {
			allNames.Insert(vol.Name)
		} else {
			allErrs = append(allErrs, el.PrefixIndex(i)...)
		}
	}
	return allNames, allErrs
}

func validateSource(source *VolumeSource) errs.ErrorList {
	numVolumes := 0
	allErrs := errs.ErrorList{}
	if source.HostDirectory != nil {
		numVolumes++
		allErrs = append(allErrs, validateHostDir(source.HostDirectory).Prefix("hostDirectory")...)
	}
	if source.EmptyDirectory != nil {
		numVolumes++
		//EmptyDirs have nothing to validate
	}
	if numVolumes != 1 {
		allErrs = append(allErrs, errs.NewInvalid("", source))
	}
	return allErrs
}

func validateHostDir(hostDir *HostDirectory) errs.ErrorList {
	allErrs := errs.ErrorList{}
	if hostDir.Path == "" {
		allErrs = append(allErrs, errs.NewNotFound("path", hostDir.Path))
	}
	return allErrs
}

var supportedPortProtocols = util.NewStringSet("TCP", "UDP")

func validatePorts(ports []Port) errs.ErrorList {
	allErrs := errs.ErrorList{}

	allNames := util.StringSet{}
	for i := range ports {
		pErrs := errs.ErrorList{}
		port := &ports[i] // so we can set default values
		if len(port.Name) > 0 {
			if len(port.Name) > 63 || !util.IsDNSLabel(port.Name) {
				pErrs = append(pErrs, errs.NewInvalid("name", port.Name))
			} else if allNames.Has(port.Name) {
				pErrs = append(pErrs, errs.NewDuplicate("name", port.Name))
			} else {
				allNames.Insert(port.Name)
			}
		}
		if port.ContainerPort == 0 {
			pErrs = append(pErrs, errs.NewRequired("containerPort", port.ContainerPort))
		} else if !util.IsValidPortNum(port.ContainerPort) {
			pErrs = append(pErrs, errs.NewInvalid("containerPort", port.ContainerPort))
		}
		if port.HostPort != 0 && !util.IsValidPortNum(port.HostPort) {
			pErrs = append(pErrs, errs.NewInvalid("hostPort", port.HostPort))
		}
		if len(port.Protocol) == 0 {
			port.Protocol = "TCP"
		} else if !supportedPortProtocols.Has(strings.ToUpper(port.Protocol)) {
			pErrs = append(pErrs, errs.NewNotSupported("protocol", port.Protocol))
		}
		allErrs = append(allErrs, pErrs.PrefixIndex(i)...)
	}
	return allErrs
}

func validateEnv(vars []EnvVar) errs.ErrorList {
	allErrs := errs.ErrorList{}

	for i := range vars {
		vErrs := errs.ErrorList{}
		ev := &vars[i] // so we can set default values
		if len(ev.Name) == 0 {
			vErrs = append(vErrs, errs.NewRequired("name", ev.Name))
		}
		if !util.IsCIdentifier(ev.Name) {
			vErrs = append(vErrs, errs.NewInvalid("name", ev.Name))
		}
		allErrs = append(allErrs, vErrs.PrefixIndex(i)...)
	}
	return allErrs
}

func validateVolumeMounts(mounts []VolumeMount, volumes util.StringSet) errs.ErrorList {
	allErrs := errs.ErrorList{}

	for i := range mounts {
		mErrs := errs.ErrorList{}
		mnt := &mounts[i] // so we can set default values
		if len(mnt.Name) == 0 {
			mErrs = append(mErrs, errs.NewRequired("name", mnt.Name))
		} else if !volumes.Has(mnt.Name) {
			mErrs = append(mErrs, errs.NewNotFound("name", mnt.Name))
		}
		if len(mnt.MountPath) == 0 {
			mErrs = append(mErrs, errs.NewRequired("mountPath", mnt.MountPath))
		}
		allErrs = append(allErrs, mErrs.PrefixIndex(i)...)
	}
	return allErrs
}

// AccumulateUniquePorts runs an extraction function on each Port of each Container,
// accumulating the results and returning an error if any ports conflict.
func AccumulateUniquePorts(containers []Container, accumulator map[int]bool, extract func(*Port) int) errs.ErrorList {
	allErrs := errs.ErrorList{}

	for ci := range containers {
		cErrs := errs.ErrorList{}
		ctr := &containers[ci]
		for pi := range ctr.Ports {
			port := extract(&ctr.Ports[pi])
			if port == 0 {
				continue
			}
			if accumulator[port] {
				cErrs = append(cErrs, errs.NewDuplicate("Port", port))
			} else {
				accumulator[port] = true
			}
		}
		allErrs = append(allErrs, cErrs.PrefixIndex(ci)...)
	}
	return allErrs
}

// Checks for colliding Port.HostPort values across a slice of containers.
func checkHostPortConflicts(containers []Container) errs.ErrorList {
	allPorts := map[int]bool{}
	return AccumulateUniquePorts(containers, allPorts, func(p *Port) int { return p.HostPort })
}

func validateContainers(containers []Container, volumes util.StringSet) errs.ErrorList {
	allErrs := errs.ErrorList{}

	allNames := util.StringSet{}
	for i := range containers {
		cErrs := errs.ErrorList{}
		ctr := &containers[i] // so we can set default values
		if len(ctr.Name) == 0 {
			cErrs = append(cErrs, errs.NewRequired("name", ctr.Name))
		} else if !util.IsDNSLabel(ctr.Name) {
			cErrs = append(cErrs, errs.NewInvalid("name", ctr.Name))
		} else if allNames.Has(ctr.Name) {
			cErrs = append(cErrs, errs.NewDuplicate("name", ctr.Name))
		} else {
			allNames.Insert(ctr.Name)
		}
		if len(ctr.Image) == 0 {
			cErrs = append(cErrs, errs.NewRequired("image", ctr.Image))
		}
		cErrs = append(cErrs, validatePorts(ctr.Ports).Prefix("ports")...)
		cErrs = append(cErrs, validateEnv(ctr.Env).Prefix("env")...)
		cErrs = append(cErrs, validateVolumeMounts(ctr.VolumeMounts, volumes).Prefix("volumeMounts")...)
		allErrs = append(allErrs, cErrs.PrefixIndex(i)...)
	}
	// Check for colliding ports across all containers.
	// TODO(thockin): This really is dependent on the network config of the host (IP per pod?)
	// and the config of the new manifest.  But we have not specced that out yet, so we'll just
	// make some assumptions for now.  As of now, pods share a network namespace, which means that
	// every Port.HostPort across the whole pod must be unique.
	allErrs = append(allErrs, checkHostPortConflicts(containers)...)

	return allErrs
}

var supportedManifestVersions = util.NewStringSet("v1beta1", "v1beta2")

// ValidateManifest tests that the specified ContainerManifest has valid data.
// This includes checking formatting and uniqueness.  It also canonicalizes the
// structure by setting default values and implementing any backwards-compatibility
// tricks.
func ValidateManifest(manifest *ContainerManifest) errs.ErrorList {
	allErrs := errs.ErrorList{}

	if len(manifest.Version) == 0 {
		allErrs = append(allErrs, errs.NewRequired("version", manifest.Version))
	} else if !supportedManifestVersions.Has(strings.ToLower(manifest.Version)) {
		allErrs = append(allErrs, errs.NewNotSupported("version", manifest.Version))
	}
	allVolumes, errs := validateVolumes(manifest.Volumes)
	allErrs = append(allErrs, errs.Prefix("volumes")...)
	allErrs = append(allErrs, validateContainers(manifest.Containers, allVolumes).Prefix("containers")...)
	return allErrs
}

func ValidatePodState(podState *PodState) errs.ErrorList {
	allErrs := errs.ErrorList(ValidateManifest(&podState.Manifest)).Prefix("manifest")
	if podState.RestartPolicy.Type == "" {
		podState.RestartPolicy.Type = RestartAlways
	} else if podState.RestartPolicy.Type != RestartAlways &&
		podState.RestartPolicy.Type != RestartOnFailure &&
		podState.RestartPolicy.Type != RestartNever {
		allErrs = append(allErrs, errs.NewNotSupported("restartPolicy.type", podState.RestartPolicy.Type))
	}

	return allErrs
}

// Pod tests if required fields in the pod are set.
func ValidatePod(pod *Pod) errs.ErrorList {
	allErrs := errs.ErrorList{}
	if len(pod.ID) == 0 {
		allErrs = append(allErrs, errs.NewRequired("id", pod.ID))
	}
	allErrs = append(allErrs, ValidatePodState(&pod.DesiredState).Prefix("desiredState")...)
	return allErrs
}

// ValidateService tests if required fields in the service are set.
func ValidateService(service *Service) errs.ErrorList {
	allErrs := errs.ErrorList{}
	if len(service.ID) == 0 {
		allErrs = append(allErrs, errs.NewRequired("id", service.ID))
	} else if !util.IsDNS952Label(service.ID) {
		allErrs = append(allErrs, errs.NewInvalid("id", service.ID))
	}
	if !util.IsValidPortNum(service.Port) {
		allErrs = append(allErrs, errs.NewInvalid("Service.Port", service.Port))
	}
	if labels.Set(service.Selector).AsSelector().Empty() {
		allErrs = append(allErrs, errs.NewRequired("selector", service.Selector))
	}
	return allErrs
}

// ValidateReplicationController tests if required fields in the replication controller are set.
func ValidateReplicationController(controller *ReplicationController) errs.ErrorList {
	allErrs := errs.ErrorList{}
	if len(controller.ID) == 0 {
		allErrs = append(allErrs, errs.NewRequired("id", controller.ID))
	}
	if labels.Set(controller.DesiredState.ReplicaSelector).AsSelector().Empty() {
		allErrs = append(allErrs, errs.NewRequired("desiredState.replicaSelector", controller.DesiredState.ReplicaSelector))
	}
	selector := labels.Set(controller.DesiredState.ReplicaSelector).AsSelector()
	labels := labels.Set(controller.DesiredState.PodTemplate.Labels)
	if !selector.Matches(labels) {
		allErrs = append(allErrs, errs.NewInvalid("desiredState.podTemplate.labels", controller.DesiredState.PodTemplate))
	}
	if controller.DesiredState.Replicas < 0 {
		allErrs = append(allErrs, errs.NewInvalid("desiredState.replicas", controller.DesiredState.Replicas))
	}
	allErrs = append(allErrs, ValidateManifest(&controller.DesiredState.PodTemplate.DesiredState.Manifest).Prefix("desiredState.podTemplate.desiredState.manifest")...)
	return allErrs
}
<|endoftext|>


<commit_before>package ps

import . "strconv"

import "testing"
import "sort"

func TestMapNil(t *testing.T) {
    m := NewMap()
    keys := m.Keys()
    if len(keys) != 0 {
        t.Errorf("Empty map has keys")
    }
}

func TestMapImmutable(t *testing.T) {
    // build a couple small maps
    world := NewMap().Set("hello", "world")
    kids := world.Set("hello", "kids")

    // both maps should still retain their data
    if v, _ := world.Lookup("hello"); v != "world" {
        t.Errorf("Set() modified the receiving map")
    }
    if size := world.Size(); size != 1 {
        t.Errorf("world size is not 1 : %d", size)
    }
    if v, _ := kids.Lookup("hello"); v != "kids" {
        t.Errorf("Set() did not modify the resulting map")
    }
    if size := kids.Size(); size != 1 {
        t.Errorf("kids size is not 1 : %d", size)
    }

    // both maps have the right keys
    if keys := world.Keys(); len(keys) != 1 || keys[0] != "hello" {
        t.Errorf("world has the wrong keys: %#v", keys)
    }
    if keys := kids.Keys(); len(keys) != 1 || keys[0] != "hello" {
        t.Errorf("kids has the wrong keys: %#v", keys)
    }

    // test deletion
    empty := kids.Delete("hello")
    if size := empty.Size(); size != 0 {
        t.Errorf("empty size is not 1 : %d", size)
    }
    if keys := empty.Keys(); len(keys) != 0 {
        t.Errorf("empty has the wrong keys: %#v", keys)
    }
}

func TestMapMultipleKeys(t *testing.T) {
    // map with multiple keys each with pointer values
    one := 1
    two := 2
    three := 3
    m := NewMap().Set("one", &one).Set("two", &two).Set("three", &three)

    // do we have the right number of keys?
    keys := m.Keys()
    if len(keys) != 3 {
        t.Logf("wrong size keys: %d", len(keys))
        t.FailNow()
    }

    // do we have the right keys?
    sort.Strings(keys)
    if keys[0] != "one" {
        t.Errorf("unexpected key: %s", keys[0])
    }
    if keys[1] != "three" {
        t.Errorf("unexpected key: %s", keys[1])
    }
    if keys[2] != "two" {
        t.Errorf("unexpected key: %s", keys[2])
    }

    // do we have the right values?
    vp, ok := m.Lookup("one");
    if !ok {
        t.Logf("missing value for one")
        t.FailNow()
    }
    if v := vp.(*int); *v != 1 {
        t.Errorf("wrong value: %d\n", *v)
    }
    vp, ok = m.Lookup("two");
    if !ok {
        t.Logf("missing value for two")
        t.FailNow()
    }
    if v := vp.(*int); *v != 2 {
        t.Errorf("wrong value: %d\n", *v)
    }
    vp, ok = m.Lookup("three");
    if !ok {
        t.Logf("missing value for three")
        t.FailNow()
    }
    if v := vp.(*int); *v != 3 {
        t.Errorf("wrong value: %d\n", *v)
    }
}

func TestMapManyKeys (t *testing.T) {
    // build a map with many keys and values
    count := 100
    m := NewMap()
    for i := 0; i<count; i++ {
        m = m.Set(Itoa(i), i)
    }

    if m.Size() != 100 {
        t.Errorf("Wrong number of keys", m.Size())
    }

    m = m.Delete("42").Delete("7").Delete("19").Delete("99")
    if m.Size() != 96 {
        t.Errorf("Wrong number of keys", m.Size())
    }

    for i:=43; i<99; i++ {
        v, ok := m.Lookup(Itoa(i))
        if !ok || v != i {
            t.Errorf("Wrong value for key %d", i)
        }
    }
}

func BenchmarkMapSet(b *testing.B) {
    m := NewMap()
    for i := 0; i < b.N; i++ {
        m = m.Set("foo", i)
    }
}

func BenchmarkMapDelete(b *testing.B) {
    m := NewMap().Set("key", "value")
    for i := 0; i < b.N; i++ {
        m.Delete("key")
    }
}
<commit_msg>Benchmark for key hashing<commit_after>package ps

import . "strconv"

import "testing"
import "sort"

func TestMapNil(t *testing.T) {
    m := NewMap()
    keys := m.Keys()
    if len(keys) != 0 {
        t.Errorf("Empty map has keys")
    }
}

func TestMapImmutable(t *testing.T) {
    // build a couple small maps
    world := NewMap().Set("hello", "world")
    kids := world.Set("hello", "kids")

    // both maps should still retain their data
    if v, _ := world.Lookup("hello"); v != "world" {
        t.Errorf("Set() modified the receiving map")
    }
    if size := world.Size(); size != 1 {
        t.Errorf("world size is not 1 : %d", size)
    }
    if v, _ := kids.Lookup("hello"); v != "kids" {
        t.Errorf("Set() did not modify the resulting map")
    }
    if size := kids.Size(); size != 1 {
        t.Errorf("kids size is not 1 : %d", size)
    }

    // both maps have the right keys
    if keys := world.Keys(); len(keys) != 1 || keys[0] != "hello" {
        t.Errorf("world has the wrong keys: %#v", keys)
    }
    if keys := kids.Keys(); len(keys) != 1 || keys[0] != "hello" {
        t.Errorf("kids has the wrong keys: %#v", keys)
    }

    // test deletion
    empty := kids.Delete("hello")
    if size := empty.Size(); size != 0 {
        t.Errorf("empty size is not 1 : %d", size)
    }
    if keys := empty.Keys(); len(keys) != 0 {
        t.Errorf("empty has the wrong keys: %#v", keys)
    }
}

func TestMapMultipleKeys(t *testing.T) {
    // map with multiple keys each with pointer values
    one := 1
    two := 2
    three := 3
    m := NewMap().Set("one", &one).Set("two", &two).Set("three", &three)

    // do we have the right number of keys?
    keys := m.Keys()
    if len(keys) != 3 {
        t.Logf("wrong size keys: %d", len(keys))
        t.FailNow()
    }

    // do we have the right keys?
    sort.Strings(keys)
    if keys[0] != "one" {
        t.Errorf("unexpected key: %s", keys[0])
    }
    if keys[1] != "three" {
        t.Errorf("unexpected key: %s", keys[1])
    }
    if keys[2] != "two" {
        t.Errorf("unexpected key: %s", keys[2])
    }

    // do we have the right values?
    vp, ok := m.Lookup("one");
    if !ok {
        t.Logf("missing value for one")
        t.FailNow()
    }
    if v := vp.(*int); *v != 1 {
        t.Errorf("wrong value: %d\n", *v)
    }
    vp, ok = m.Lookup("two");
    if !ok {
        t.Logf("missing value for two")
        t.FailNow()
    }
    if v := vp.(*int); *v != 2 {
        t.Errorf("wrong value: %d\n", *v)
    }
    vp, ok = m.Lookup("three");
    if !ok {
        t.Logf("missing value for three")
        t.FailNow()
    }
    if v := vp.(*int); *v != 3 {
        t.Errorf("wrong value: %d\n", *v)
    }
}

func TestMapManyKeys (t *testing.T) {
    // build a map with many keys and values
    count := 100
    m := NewMap()
    for i := 0; i<count; i++ {
        m = m.Set(Itoa(i), i)
    }

    if m.Size() != 100 {
        t.Errorf("Wrong number of keys", m.Size())
    }

    m = m.Delete("42").Delete("7").Delete("19").Delete("99")
    if m.Size() != 96 {
        t.Errorf("Wrong number of keys", m.Size())
    }

    for i:=43; i<99; i++ {
        v, ok := m.Lookup(Itoa(i))
        if !ok || v != i {
            t.Errorf("Wrong value for key %d", i)
        }
    }
}

func BenchmarkMapSet(b *testing.B) {
    m := NewMap()
    for i := 0; i < b.N; i++ {
        m = m.Set("foo", i)
    }
}

func BenchmarkMapDelete(b *testing.B) {
    m := NewMap().Set("key", "value")
    for i := 0; i < b.N; i++ {
        m.Delete("key")
    }
}

func BenchmarkHashKey(b *testing.B) {
    key := "this is a key"
    for i := 0; i < b.N; i++ {
        _ = hashKey(key)
    }
}
<|endoftext|>


<commit_before>package compiler

import (
	"bytes"
	"fmt"
	"reflect"
	"strings"

	"github.com/Zac-Garby/pluto/ast"
	"github.com/Zac-Garby/pluto/bytecode"
	"github.com/Zac-Garby/pluto/object"
)

// CompileExpression compiles an AST expression.
func (c *Compiler) CompileExpression(n ast.Expression) error {
	switch node := n.(type) {
	case *ast.InfixExpression:
		return c.compileInfix(node)
	case *ast.PrefixExpression:
		return c.compilePrefix(node)
	case *ast.Number:
		return c.compileNumber(node)
	case *ast.String:
		return c.compileString(node)
	case *ast.Boolean:
		return c.compileBoolean(node)
	case *ast.Char:
		return c.compileChar(node)
	case *ast.Null:
		return c.compileNull(node)
	case *ast.Identifier:
		return c.compileIdentifier(node)
	case *ast.Array:
		return c.compileArray(node)
	case *ast.Tuple:
		return c.compileTuple(node)
	case *ast.Map:
		return c.compileMap(node)
	case *ast.AssignExpression:
		return c.compileAssign(node)
	case *ast.IfExpression:
		return c.compileIf(node)
	case *ast.WhileLoop:
		return c.compileWhile(node)
	case *ast.FunctionCall:
		return c.compileFnCall(node)
	case *ast.Argument:
		return c.CompileExpression(node.Value)
	default:
		return fmt.Errorf("compiler: compilation not yet implemented for %s", reflect.TypeOf(n))
	}
}

func (c *Compiler) compileNumber(node *ast.Number) error {
	obj := &object.Number{Value: node.Value}
	c.Constants = append(c.Constants, obj)
	index := len(c.Constants) - 1

	if index >= 1<<16 {
		return fmt.Errorf("compiler: constant index %d greater than 1 << 16 (maximum uint16)", index)
	}

	low, high := runeToBytes(rune(index))

	c.Bytes = append(c.Bytes, bytecode.LoadConst, high, low)

	return nil
}

func (c *Compiler) compileString(node *ast.String) error {
	obj := &object.String{Value: node.Value}
	c.Constants = append(c.Constants, obj)
	index := len(c.Constants) - 1

	if index >= 1<<16 {
		return fmt.Errorf("compiler: constant index %d greater than 1 << 16 (maximum uint16)", index)
	}

	low, high := runeToBytes(rune(index))

	c.Bytes = append(c.Bytes, bytecode.LoadConst, high, low)

	return nil
}

func (c *Compiler) compileBoolean(node *ast.Boolean) error {
	obj := &object.Boolean{Value: node.Value}
	c.Constants = append(c.Constants, obj)
	index := len(c.Constants) - 1

	if index >= 1<<16 {
		return fmt.Errorf("compiler: constant index %d greater than 1 << 16 (maximum uint16)", index)
	}

	low, high := runeToBytes(rune(index))

	c.Bytes = append(c.Bytes, bytecode.LoadConst, high, low)

	return nil
}

func (c *Compiler) compileChar(node *ast.Char) error {
	obj := &object.Char{Value: rune(node.Value)}
	c.Constants = append(c.Constants, obj)
	index := len(c.Constants) - 1

	if index >= 1<<16 {
		return fmt.Errorf("compiler: constant index %d greater than 1 << 16 (maximum uint16)", index)
	}

	low, high := runeToBytes(rune(index))

	c.Bytes = append(c.Bytes, bytecode.LoadConst, high, low)

	return nil
}

func (c *Compiler) compileNull(node *ast.Null) error {
	obj := object.NullObj
	c.Constants = append(c.Constants, obj)
	index := len(c.Constants) - 1

	if index >= 1<<16 {
		return fmt.Errorf("compiler: constant index %d greater than 1 << 16 (maximum uint16)", index)
	}

	low, high := runeToBytes(rune(index))

	c.Bytes = append(c.Bytes, bytecode.LoadConst, high, low)

	return nil
}

func (c *Compiler) compileIdentifier(node *ast.Identifier) error {
	var index int

	for i, name := range c.Names {
		if name == node.Value {
			index = i
			goto found
		}
	}

	// These two lines are executed if the name isn't found
	c.Names = append(c.Names, node.Value)
	index = len(c.Names) - 1

found:
	low, high := runeToBytes(rune(index))

	c.Bytes = append(c.Bytes, bytecode.LoadName, high, low)

	return nil
}

func (c *Compiler) compileAssign(node *ast.AssignExpression) error {
	if err := c.CompileExpression(node.Value); err != nil {
		return err
	}

	c.Names = append(c.Names, node.Name.(*ast.Identifier).Value)
	index := len(c.Names) - 1

	if index >= 1<<16 {
		return fmt.Errorf("compiler: name index %d greater than 1 << 16 (maximum uint16)", index)
	}

	low, high := runeToBytes(rune(index))

	c.Bytes = append(c.Bytes, bytecode.StoreName, high, low)

	return nil
}

func (c *Compiler) compileInfix(node *ast.InfixExpression) error {
	left, right := node.Left, node.Right

	if err := c.CompileExpression(left); err != nil {
		return err
	}

	if err := c.CompileExpression(right); err != nil {
		return err
	}

	op, ok := map[string]byte{
		"+":  bytecode.BinaryAdd,
		"-":  bytecode.BinarySubtract,
		"*":  bytecode.BinaryMultiply,
		"/":  bytecode.BinaryDivide,
		"**": bytecode.BinaryExponent,
		"//": bytecode.BinaryFloorDiv,
		"%":  bytecode.BinaryFloorDiv,
		"||": bytecode.BinaryOr,
		"&&": bytecode.BinaryAnd,
		"|":  bytecode.BinaryBitOr,
		"&":  bytecode.BinaryBitAnd,
		"==": bytecode.BinaryEquals,
		"!=": bytecode.BinaryNotEqual,
		"<":  bytecode.BinaryLessThan,
		">":  bytecode.BinaryMoreThan,
		"<=": bytecode.BinaryLessEq,
		">=": bytecode.BinaryMoreEq,
	}[node.Operator]

	if !ok {
		return fmt.Errorf("compiler: operator %s not yet implemented", node.Operator)
	}

	c.Bytes = append(c.Bytes, op)

	return nil
}

func (c *Compiler) compilePrefix(node *ast.PrefixExpression) error {
	if err := c.CompileExpression(node.Right); err != nil {
		return err
	}

	op := map[string]byte{
		"+": bytecode.UnaryNoOp,
		"-": bytecode.UnaryNegate,
		"!": bytecode.UnaryInvert,
	}[node.Operator]

	c.Bytes = append(c.Bytes, op)

	return nil
}

func (c *Compiler) compileIf(node *ast.IfExpression) error {
	if err := c.CompileExpression(node.Condition); err != nil {
		return err
	}

	// JumpIfFalse (82) with 2 empty argument bytes
	c.Bytes = append(c.Bytes, bytecode.JumpIfFalse, 0, 0)
	condJump := len(c.Bytes) - 3

	if err := c.CompileStatement(node.Consequence); err != nil {
		return err
	}

	var skipJump int

	if node.Alternative != nil {
		// Jump past the alternative
		c.Bytes = append(c.Bytes, bytecode.Jump, 0, 0)
		skipJump = len(c.Bytes) - 3
	}

	// Set the jump target after the conditional
	condIndex := rune(len(c.Bytes))
	low, high := runeToBytes(condIndex)
	c.Bytes[condJump+1] = high
	c.Bytes[condJump+2] = low

	if node.Alternative != nil {
		if err := c.CompileStatement(node.Alternative); err != nil {
			return err
		}

		// Set the jump target after the conditional
		skipIndex := rune(len(c.Bytes))
		low, high = runeToBytes(skipIndex)
		c.Bytes[skipJump+1] = high
		c.Bytes[skipJump+2] = low
	}

	return nil
}

func (c *Compiler) compileArray(node *ast.Array) error {
	for _, elem := range node.Elements {
		if err := c.CompileExpression(elem); err != nil {
			return err
		}
	}

	low, high := runeToBytes(rune(len(node.Elements)))

	c.Bytes = append(c.Bytes, bytecode.MakeArray, high, low)

	return nil
}

func (c *Compiler) compileTuple(node *ast.Tuple) error {
	for _, elem := range node.Value {
		if err := c.CompileExpression(elem); err != nil {
			return err
		}
	}

	low, high := runeToBytes(rune(len(node.Value)))

	c.Bytes = append(c.Bytes, bytecode.MakeTuple, high, low)

	return nil
}

func (c *Compiler) compileMap(node *ast.Map) error {
	for key, val := range node.Pairs {
		if err := c.CompileExpression(key); err != nil {
			return err
		}

		if err := c.CompileExpression(val); err != nil {
			return err
		}
	}

	low, high := runeToBytes(rune(len(node.Pairs)))

	c.Bytes = append(c.Bytes, bytecode.MakeMap, high, low)

	return nil
}

func (c *Compiler) compileWhile(node *ast.WhileLoop) error {
	// Jump here to go to the next iteration
	start := len(c.Bytes) - 1

	if err := c.CompileExpression(node.Condition); err != nil {
		return err
	}

	// An empty jump to the end of the loop
	c.Bytes = append(c.Bytes, bytecode.JumpIfFalse, 0, 0)
	skipJump := len(c.Bytes) - 3

	// Compile the loop's body
	if err := c.CompileStatement(node.Body); err != nil {
		return err
	}

	// After the body, jump back to the beginning of the loop
	low, high := runeToBytes(rune(start))
	c.Bytes = append(c.Bytes, bytecode.Jump, high, low)

	// If the condition isn't met, jump to the end of the loop
	skipIndex := rune(len(c.Bytes))
	low, high = runeToBytes(skipIndex)
	c.Bytes[skipJump+1] = high
	c.Bytes[skipJump+2] = low

	return nil
}

func (c *Compiler) compileFnCall(node *ast.FunctionCall) error {
	var pstring bytes.Buffer

	for _, item := range node.Pattern {
		if id, ok := item.(*ast.Identifier); ok {
			pstring.WriteString(id.Value + " ")
		} else {
			pstring.WriteString("$ ")
		}
	}

	str := strings.TrimSpace(pstring.String())
	c.Patterns = append(c.Patterns, str)

	for _, item := range node.Pattern {
		if arg, ok := item.(*ast.Argument); ok {
			if err := c.CompileExpression(arg); err != nil {
				return err
			}
		}
	}

	low, high := runeToBytes(rune(len(c.Patterns) - 1))
	c.Bytes = append(c.Bytes, bytecode.Call, high, low)

	return nil
}
<commit_msg>Improve pattern-string creation<commit_after>package compiler

import (
	"fmt"
	"reflect"
	"strings"

	"github.com/Zac-Garby/pluto/ast"
	"github.com/Zac-Garby/pluto/bytecode"
	"github.com/Zac-Garby/pluto/object"
)

// CompileExpression compiles an AST expression.
func (c *Compiler) CompileExpression(n ast.Expression) error {
	switch node := n.(type) {
	case *ast.InfixExpression:
		return c.compileInfix(node)
	case *ast.PrefixExpression:
		return c.compilePrefix(node)
	case *ast.Number:
		return c.compileNumber(node)
	case *ast.String:
		return c.compileString(node)
	case *ast.Boolean:
		return c.compileBoolean(node)
	case *ast.Char:
		return c.compileChar(node)
	case *ast.Null:
		return c.compileNull(node)
	case *ast.Identifier:
		return c.compileIdentifier(node)
	case *ast.Array:
		return c.compileArray(node)
	case *ast.Tuple:
		return c.compileTuple(node)
	case *ast.Map:
		return c.compileMap(node)
	case *ast.AssignExpression:
		return c.compileAssign(node)
	case *ast.IfExpression:
		return c.compileIf(node)
	case *ast.WhileLoop:
		return c.compileWhile(node)
	case *ast.FunctionCall:
		return c.compileFnCall(node)
	case *ast.Argument:
		return c.CompileExpression(node.Value)
	default:
		return fmt.Errorf("compiler: compilation not yet implemented for %s", reflect.TypeOf(n))
	}
}

func (c *Compiler) compileNumber(node *ast.Number) error {
	obj := &object.Number{Value: node.Value}
	c.Constants = append(c.Constants, obj)
	index := len(c.Constants) - 1

	if index >= 1<<16 {
		return fmt.Errorf("compiler: constant index %d greater than 1 << 16 (maximum uint16)", index)
	}

	low, high := runeToBytes(rune(index))

	c.Bytes = append(c.Bytes, bytecode.LoadConst, high, low)

	return nil
}

func (c *Compiler) compileString(node *ast.String) error {
	obj := &object.String{Value: node.Value}
	c.Constants = append(c.Constants, obj)
	index := len(c.Constants) - 1

	if index >= 1<<16 {
		return fmt.Errorf("compiler: constant index %d greater than 1 << 16 (maximum uint16)", index)
	}

	low, high := runeToBytes(rune(index))

	c.Bytes = append(c.Bytes, bytecode.LoadConst, high, low)

	return nil
}

func (c *Compiler) compileBoolean(node *ast.Boolean) error {
	obj := &object.Boolean{Value: node.Value}
	c.Constants = append(c.Constants, obj)
	index := len(c.Constants) - 1

	if index >= 1<<16 {
		return fmt.Errorf("compiler: constant index %d greater than 1 << 16 (maximum uint16)", index)
	}

	low, high := runeToBytes(rune(index))

	c.Bytes = append(c.Bytes, bytecode.LoadConst, high, low)

	return nil
}

func (c *Compiler) compileChar(node *ast.Char) error {
	obj := &object.Char{Value: rune(node.Value)}
	c.Constants = append(c.Constants, obj)
	index := len(c.Constants) - 1

	if index >= 1<<16 {
		return fmt.Errorf("compiler: constant index %d greater than 1 << 16 (maximum uint16)", index)
	}

	low, high := runeToBytes(rune(index))

	c.Bytes = append(c.Bytes, bytecode.LoadConst, high, low)

	return nil
}

func (c *Compiler) compileNull(node *ast.Null) error {
	obj := object.NullObj
	c.Constants = append(c.Constants, obj)
	index := len(c.Constants) - 1

	if index >= 1<<16 {
		return fmt.Errorf("compiler: constant index %d greater than 1 << 16 (maximum uint16)", index)
	}

	low, high := runeToBytes(rune(index))

	c.Bytes = append(c.Bytes, bytecode.LoadConst, high, low)

	return nil
}

func (c *Compiler) compileIdentifier(node *ast.Identifier) error {
	var index int

	for i, name := range c.Names {
		if name == node.Value {
			index = i
			goto found
		}
	}

	// These two lines are executed if the name isn't found
	c.Names = append(c.Names, node.Value)
	index = len(c.Names) - 1

found:
	low, high := runeToBytes(rune(index))

	c.Bytes = append(c.Bytes, bytecode.LoadName, high, low)

	return nil
}

func (c *Compiler) compileAssign(node *ast.AssignExpression) error {
	if err := c.CompileExpression(node.Value); err != nil {
		return err
	}

	c.Names = append(c.Names, node.Name.(*ast.Identifier).Value)
	index := len(c.Names) - 1

	if index >= 1<<16 {
		return fmt.Errorf("compiler: name index %d greater than 1 << 16 (maximum uint16)", index)
	}

	low, high := runeToBytes(rune(index))

	c.Bytes = append(c.Bytes, bytecode.StoreName, high, low)

	return nil
}

func (c *Compiler) compileInfix(node *ast.InfixExpression) error {
	left, right := node.Left, node.Right

	if err := c.CompileExpression(left); err != nil {
		return err
	}

	if err := c.CompileExpression(right); err != nil {
		return err
	}

	op, ok := map[string]byte{
		"+":  bytecode.BinaryAdd,
		"-":  bytecode.BinarySubtract,
		"*":  bytecode.BinaryMultiply,
		"/":  bytecode.BinaryDivide,
		"**": bytecode.BinaryExponent,
		"//": bytecode.BinaryFloorDiv,
		"%":  bytecode.BinaryFloorDiv,
		"||": bytecode.BinaryOr,
		"&&": bytecode.BinaryAnd,
		"|":  bytecode.BinaryBitOr,
		"&":  bytecode.BinaryBitAnd,
		"==": bytecode.BinaryEquals,
		"!=": bytecode.BinaryNotEqual,
		"<":  bytecode.BinaryLessThan,
		">":  bytecode.BinaryMoreThan,
		"<=": bytecode.BinaryLessEq,
		">=": bytecode.BinaryMoreEq,
	}[node.Operator]

	if !ok {
		return fmt.Errorf("compiler: operator %s not yet implemented", node.Operator)
	}

	c.Bytes = append(c.Bytes, op)

	return nil
}

func (c *Compiler) compilePrefix(node *ast.PrefixExpression) error {
	if err := c.CompileExpression(node.Right); err != nil {
		return err
	}

	op := map[string]byte{
		"+": bytecode.UnaryNoOp,
		"-": bytecode.UnaryNegate,
		"!": bytecode.UnaryInvert,
	}[node.Operator]

	c.Bytes = append(c.Bytes, op)

	return nil
}

func (c *Compiler) compileIf(node *ast.IfExpression) error {
	if err := c.CompileExpression(node.Condition); err != nil {
		return err
	}

	// JumpIfFalse (82) with 2 empty argument bytes
	c.Bytes = append(c.Bytes, bytecode.JumpIfFalse, 0, 0)
	condJump := len(c.Bytes) - 3

	if err := c.CompileStatement(node.Consequence); err != nil {
		return err
	}

	var skipJump int

	if node.Alternative != nil {
		// Jump past the alternative
		c.Bytes = append(c.Bytes, bytecode.Jump, 0, 0)
		skipJump = len(c.Bytes) - 3
	}

	// Set the jump target after the conditional
	condIndex := rune(len(c.Bytes))
	low, high := runeToBytes(condIndex)
	c.Bytes[condJump+1] = high
	c.Bytes[condJump+2] = low

	if node.Alternative != nil {
		if err := c.CompileStatement(node.Alternative); err != nil {
			return err
		}

		// Set the jump target after the conditional
		skipIndex := rune(len(c.Bytes))
		low, high = runeToBytes(skipIndex)
		c.Bytes[skipJump+1] = high
		c.Bytes[skipJump+2] = low
	}

	return nil
}

func (c *Compiler) compileArray(node *ast.Array) error {
	for _, elem := range node.Elements {
		if err := c.CompileExpression(elem); err != nil {
			return err
		}
	}

	low, high := runeToBytes(rune(len(node.Elements)))

	c.Bytes = append(c.Bytes, bytecode.MakeArray, high, low)

	return nil
}

func (c *Compiler) compileTuple(node *ast.Tuple) error {
	for _, elem := range node.Value {
		if err := c.CompileExpression(elem); err != nil {
			return err
		}
	}

	low, high := runeToBytes(rune(len(node.Value)))

	c.Bytes = append(c.Bytes, bytecode.MakeTuple, high, low)

	return nil
}

func (c *Compiler) compileMap(node *ast.Map) error {
	for key, val := range node.Pairs {
		if err := c.CompileExpression(key); err != nil {
			return err
		}

		if err := c.CompileExpression(val); err != nil {
			return err
		}
	}

	low, high := runeToBytes(rune(len(node.Pairs)))

	c.Bytes = append(c.Bytes, bytecode.MakeMap, high, low)

	return nil
}

func (c *Compiler) compileWhile(node *ast.WhileLoop) error {
	// Jump here to go to the next iteration
	start := len(c.Bytes) - 1

	if err := c.CompileExpression(node.Condition); err != nil {
		return err
	}

	// An empty jump to the end of the loop
	c.Bytes = append(c.Bytes, bytecode.JumpIfFalse, 0, 0)
	skipJump := len(c.Bytes) - 3

	// Compile the loop's body
	if err := c.CompileStatement(node.Body); err != nil {
		return err
	}

	// After the body, jump back to the beginning of the loop
	low, high := runeToBytes(rune(start))
	c.Bytes = append(c.Bytes, bytecode.Jump, high, low)

	// If the condition isn't met, jump to the end of the loop
	skipIndex := rune(len(c.Bytes))
	low, high = runeToBytes(skipIndex)
	c.Bytes[skipJump+1] = high
	c.Bytes[skipJump+2] = low

	return nil
}

func (c *Compiler) compileFnCall(node *ast.FunctionCall) error {
	var ptn []string

	for _, item := range node.Pattern {
		if id, ok := item.(*ast.Identifier); ok {
			ptn = append(ptn, id.Value)
		} else {
			ptn = append(ptn, "$")
		}
	}

	str := strings.Join(ptn, " ")
	c.Patterns = append(c.Patterns, str)

	for _, item := range node.Pattern {
		if arg, ok := item.(*ast.Argument); ok {
			if err := c.CompileExpression(arg); err != nil {
				return err
			}
		}
	}

	low, high := runeToBytes(rune(len(c.Patterns) - 1))
	c.Bytes = append(c.Bytes, bytecode.Call, high, low)

	return nil
}
<|endoftext|>


<commit_before>package cache

import (
	"time"

	"github.com/coredns/coredns/plugin/cache/freq"
	"github.com/coredns/coredns/plugin/pkg/response"
	"github.com/miekg/dns"
)

type item struct {
	Rcode              int
	Authoritative      bool
	AuthenticatedData  bool
	RecursionAvailable bool
	Answer             []dns.RR
	Ns                 []dns.RR
	Extra              []dns.RR

	origTTL uint32
	stored  time.Time

	*freq.Freq
}

func newItem(m *dns.Msg, now time.Time, d time.Duration) *item {
	i := new(item)
	i.Rcode = m.Rcode
	i.Authoritative = m.Authoritative
	i.AuthenticatedData = m.AuthenticatedData
	i.RecursionAvailable = m.RecursionAvailable
	i.Answer = m.Answer
	i.Ns = m.Ns
	i.Extra = make([]dns.RR, len(m.Extra))
	// Don't copy OPT record as these are hop-by-hop.
	j := 0
	for _, e := range m.Extra {
		if e.Header().Rrtype == dns.TypeOPT {
			continue
		}
		i.Extra[j] = e
		j++
	}
	i.Extra = i.Extra[:j]

	i.origTTL = uint32(d.Seconds())
	i.stored = now.UTC()

	i.Freq = new(freq.Freq)

	return i
}

// toMsg turns i into a message, it tailors the reply to m.
// The Authoritative bit is always set to 0, because the answer is from the cache.
func (i *item) toMsg(m *dns.Msg, now time.Time) *dns.Msg {
	m1 := new(dns.Msg)
	m1.SetReply(m)

	m1.Authoritative = false
	m1.AuthenticatedData = i.AuthenticatedData
	m1.RecursionAvailable = i.RecursionAvailable
	m1.Rcode = i.Rcode
	m1.Compress = true

	m1.Answer = make([]dns.RR, len(i.Answer))
	m1.Ns = make([]dns.RR, len(i.Ns))
	m1.Extra = make([]dns.RR, len(i.Extra))

	ttl := uint32(i.ttl(now))
	for j, r := range i.Answer {
		m1.Answer[j] = dns.Copy(r)
		m1.Answer[j].Header().Ttl = ttl
	}
	for j, r := range i.Ns {
		m1.Ns[j] = dns.Copy(r)
		m1.Ns[j].Header().Ttl = ttl
	}
	for j, r := range i.Extra {
		m1.Extra[j] = dns.Copy(r)
		if m1.Extra[j].Header().Rrtype != dns.TypeOPT {
			m1.Extra[j].Header().Ttl = ttl
		}
	}
	return m1
}

func (i *item) ttl(now time.Time) int {
	ttl := int(i.origTTL) - int(now.UTC().Sub(i.stored).Seconds())
	return ttl
}

func minMsgTTL(m *dns.Msg, mt response.Type) time.Duration {
	if mt != response.NoError && mt != response.NameError && mt != response.NoData {
		return 0
	}

	// No data to examine, return a short ttl as a fail safe.
	if len(m.Answer)+len(m.Ns) == 0 {
		return failSafeTTL
	}

	minTTL := maxTTL
	for _, r := range append(m.Answer, m.Ns...) {
		switch mt {
		case response.NameError, response.NoData:
			if r.Header().Rrtype == dns.TypeSOA {
				return time.Duration(r.(*dns.SOA).Minttl) * time.Second
			}
		case response.NoError, response.Delegation:
			if r.Header().Ttl < uint32(minTTL.Seconds()) {
				minTTL = time.Duration(r.Header().Ttl) * time.Second
			}
		}
	}
	return minTTL
}
<commit_msg>incl addtl rrs when computing cache ttl (#1549)<commit_after>package cache

import (
	"time"

	"github.com/coredns/coredns/plugin/cache/freq"
	"github.com/coredns/coredns/plugin/pkg/response"
	"github.com/miekg/dns"
)

type item struct {
	Rcode              int
	Authoritative      bool
	AuthenticatedData  bool
	RecursionAvailable bool
	Answer             []dns.RR
	Ns                 []dns.RR
	Extra              []dns.RR

	origTTL uint32
	stored  time.Time

	*freq.Freq
}

func newItem(m *dns.Msg, now time.Time, d time.Duration) *item {
	i := new(item)
	i.Rcode = m.Rcode
	i.Authoritative = m.Authoritative
	i.AuthenticatedData = m.AuthenticatedData
	i.RecursionAvailable = m.RecursionAvailable
	i.Answer = m.Answer
	i.Ns = m.Ns
	i.Extra = make([]dns.RR, len(m.Extra))
	// Don't copy OPT record as these are hop-by-hop.
	j := 0
	for _, e := range m.Extra {
		if e.Header().Rrtype == dns.TypeOPT {
			continue
		}
		i.Extra[j] = e
		j++
	}
	i.Extra = i.Extra[:j]

	i.origTTL = uint32(d.Seconds())
	i.stored = now.UTC()

	i.Freq = new(freq.Freq)

	return i
}

// toMsg turns i into a message, it tailors the reply to m.
// The Authoritative bit is always set to 0, because the answer is from the cache.
func (i *item) toMsg(m *dns.Msg, now time.Time) *dns.Msg {
	m1 := new(dns.Msg)
	m1.SetReply(m)

	m1.Authoritative = false
	m1.AuthenticatedData = i.AuthenticatedData
	m1.RecursionAvailable = i.RecursionAvailable
	m1.Rcode = i.Rcode
	m1.Compress = true

	m1.Answer = make([]dns.RR, len(i.Answer))
	m1.Ns = make([]dns.RR, len(i.Ns))
	m1.Extra = make([]dns.RR, len(i.Extra))

	ttl := uint32(i.ttl(now))
	for j, r := range i.Answer {
		m1.Answer[j] = dns.Copy(r)
		m1.Answer[j].Header().Ttl = ttl
	}
	for j, r := range i.Ns {
		m1.Ns[j] = dns.Copy(r)
		m1.Ns[j].Header().Ttl = ttl
	}
	for j, r := range i.Extra {
		m1.Extra[j] = dns.Copy(r)
		if m1.Extra[j].Header().Rrtype != dns.TypeOPT {
			m1.Extra[j].Header().Ttl = ttl
		}
	}
	return m1
}

func (i *item) ttl(now time.Time) int {
	ttl := int(i.origTTL) - int(now.UTC().Sub(i.stored).Seconds())
	return ttl
}

func minMsgTTL(m *dns.Msg, mt response.Type) time.Duration {
	if mt != response.NoError && mt != response.NameError && mt != response.NoData {
		return 0
	}

	// No data to examine, return a short ttl as a fail safe.
	if len(m.Answer)+len(m.Ns)+len(m.Extra) == 0 {
		return failSafeTTL
	}

	minTTL := maxTTL
	for _, r := range append(append(m.Answer, m.Ns...), m.Extra...) {
		if r.Header().Rrtype == dns.TypeOPT {
			// OPT records use TTL field for extended rcode and flags
			continue
		}
		switch mt {
		case response.NameError, response.NoData:
			if r.Header().Rrtype == dns.TypeSOA {
				return time.Duration(r.(*dns.SOA).Minttl) * time.Second
			}
		case response.NoError, response.Delegation:
			if r.Header().Ttl < uint32(minTTL.Seconds()) {
				minTTL = time.Duration(r.Header().Ttl) * time.Second
			}
		}
	}
	return minTTL
}
<|endoftext|>