Skip to content

Commit

Permalink
feat: Integrate slither
Browse files Browse the repository at this point in the history
  • Loading branch information
damilolaedwards committed Jul 23, 2024
1 parent 7966077 commit 6edd87c
Show file tree
Hide file tree
Showing 14 changed files with 395 additions and 162 deletions.
4 changes: 4 additions & 0 deletions .idea/assistant.iml

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

99 changes: 99 additions & 0 deletions api/api.go
Original file line number Diff line number Diff line change
@@ -0,0 +1,99 @@
package api

import (
"assistant/config"
"assistant/logging"
"assistant/types"
"encoding/json"
"fmt"
"github.com/gorilla/mux"
"github.com/rs/zerolog"
"net"
"net/http"
"os"
"os/signal"
"syscall"
)

func Start(config *config.ProjectConfig, slitherOutput *types.SlitherOutput) {
port := fmt.Sprint(":", config.Port)

if port == ":" {
port = ":8080" // Default port
}

// Create sub-logger for api module
logger := logging.NewLogger(zerolog.InfoLevel)
logger.AddWriter(os.Stdout, logging.UNSTRUCTURED, true)

// Create a new router
router := mux.NewRouter()

// Attach middleware
attachMiddleware(router)

// Serve the contracts on a subrouter
router.HandleFunc("/contracts", func(w http.ResponseWriter, r *http.Request) {
w.WriteHeader(http.StatusOK)
err := json.NewEncoder(w).Encode(slitherOutput)
if err != nil {
logger.Error("Failed to encode contracts: ", err)
return
}
})

var listener net.Listener
var err error

for i := 0; i < 10; i++ {
listener, err = net.Listen("tcp", port)
if err == nil {
break
}

logger.Info("Server failed to start on port ", port[1:])
port = incrementPort(port)
}

// Stop further execution if the server failed to start
if listener == nil {
logger.Error("Failed to start server: ", err)
return
}

logger.Info("Server started on port ", port[1:])

// Create a channel to receive interrupt signals
sigChan := make(chan os.Signal, 1)
signal.Notify(sigChan, os.Interrupt, syscall.SIGTERM)

// Start the server in a separate goroutine
serverErrorChan := make(chan error, 1)
go func() {
serverErrorChan <- http.Serve(listener, router)
}()

// Gracefully shutdown the server if a server error is encountered
select {
case <-sigChan:
logger.Info("Shutting down server...")
err := listener.Close()
if err != nil {
logger.Error("Failed to shut down server: ", err)
return
}
case err := <-serverErrorChan:
logger.Error("Server error: ", err)
}
}

func incrementPort(port string) string {
var portNum int

_, err := fmt.Sscanf(port, ":%d", &portNum)
if err != nil {
panic(err)
}

return fmt.Sprintf(":%d", portNum+1)
}
24 changes: 24 additions & 0 deletions api/middleware.go
Original file line number Diff line number Diff line change
@@ -0,0 +1,24 @@
package api

import (
"github.com/gorilla/mux"
"net/http"
)

func setHeaders(next http.Handler) http.Handler {
return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
// Set default headers
w.Header().Set("Content-Type", "application/json")

// Handle CORS headers
w.Header().Set("Access-Control-Allow-Origin", "*")
w.Header().Set("Access-Control-Allow-Methods", "POST, GET, OPTIONS, PUT, DELETE")
w.Header().Set("Access-Control-Allow-Headers", "Accept, Content-Type, Content-Length, Accept-Encoding, X-CSRF-Token, Authorization")

next.ServeHTTP(w, r)
})
}

func attachMiddleware(router *mux.Router) {
router.Use(setHeaders)
}
97 changes: 11 additions & 86 deletions cmd/generate.go
Original file line number Diff line number Diff line change
@@ -1,15 +1,14 @@
package cmd

import (
"assistant/api"
"assistant/config"
"assistant/llm"
"assistant/logging/colors"
"assistant/utils"
"assistant/slither"
"fmt"
"github.com/spf13/cobra"
"os"
"path/filepath"
"time"
)

var generateCmd = &cobra.Command{
Expand Down Expand Up @@ -111,104 +110,30 @@ func cmdRunGenerate(cmd *cobra.Command, args []string) error {
return err
}

targetContracts, err := utils.ReadDirectoryContents(projectConfig.TargetContracts.Dir, projectConfig.TargetContracts.ExcludePaths...)
// Run slither
cmdLogger.Info("Running Slither on the target contracts directory: ", colors.Green, projectConfig.TargetContracts.Dir, colors.Reset, ", Excluding paths: ", colors.Red, projectConfig.TargetContracts.ExcludePaths, colors.Reset, "\n")
slitherOutput, err := slither.RunSlitherOnDir(projectConfig.TargetContracts.Dir, projectConfig.TargetContracts.ExcludePaths...)
if err != nil {
cmdLogger.Error("Failed to run the generate command", err)
return err
}
cmdLogger.Info("Successfully ran Slither on the target contracts directory")

var fuzzTests string
if projectConfig.FuzzTests.Dir != "" {
fuzzTests, err = utils.ReadDirectoryContents(projectConfig.FuzzTests.Dir, projectConfig.FuzzTests.ExcludePaths...)
if err != nil {
cmdLogger.Error("Failed to run the generate command", err)
return err
}
}

var unitTests string
if projectConfig.UnitTests.Dir != "" {
unitTests, err = utils.ReadDirectoryContents(projectConfig.UnitTests.Dir, projectConfig.UnitTests.ExcludePaths...)
if err != nil {
cmdLogger.Error("Failed to run the generate command", err)
return err
}
}

var parsedCoverageReport utils.CoverageReport
if projectConfig.CoverageReportFile != "" {
_, err := os.Stat(projectConfig.CoverageReportFile)
if err != nil {
cmdLogger.Error("Failed to run the generate command", err)
return err
}

coverageReport, err := os.ReadFile(projectConfig.CoverageReportFile)
if err != nil {
cmdLogger.Error("Failed to run the generate command", err)
return err
}

parsedCoverageReport, err = utils.ParseCoverageReportHTML(string(coverageReport))
if err != nil {
cmdLogger.Error("Failed to run the generate command", err)
return err
}

// Exclude reports of files we do not need
includePaths := []string{projectConfig.TargetContracts.Dir}
excludePaths := projectConfig.TargetContracts.ExcludePaths

if projectConfig.FuzzTests.Dir != "" {
includePaths = append(includePaths, projectConfig.FuzzTests.Dir)
excludePaths = append(excludePaths, projectConfig.FuzzTests.ExcludePaths...)
}

utils.FilterCoverageFiles(&parsedCoverageReport, includePaths, excludePaths)
}

//var supportingFileIds []string
//if len(projectConfig.SupportingFiles) > 0 {
// supportingFileIds, err = llm.UploadFilesToOpenAI(projectConfig.SupportingFiles)
// if err != nil {
// cmdLogger.Error("Failed to run the generate command", err)
// return err
// }
//}

// Construct the messages for the LLM
messages := append(llm.TrainingPrompts(), llm.Message{
Role: "user",
Content: llm.GenerateInvariantsPrompt(projectConfig.NumInvariants, targetContracts, fuzzTests, unitTests, fmt.Sprintf("%v", parsedCoverageReport)),
})

// Make request to LLM to generate invariants
invariants, err := llm.AskGPT4Turbo(messages)
if err != nil {
cmdLogger.Error("Failed to run the generate command", err)
return err
}

invariants = fmt.Sprintf("=============== Invariants generated at %v ===============\n\n%v\n\n", time.Now().String(), invariants)

// Open the out file in append mode, create it if it doesn't exist
file, err := os.OpenFile(projectConfig.Out, os.O_APPEND|os.O_CREATE|os.O_WRONLY, 0644)
// Write contracts to a file
file, err := os.OpenFile("contracts.json", os.O_APPEND|os.O_CREATE|os.O_WRONLY, 0644)
if err != nil {
return err
}

defer func(file *os.File) {
err := file.Close()
if err != nil {
cmdLogger.Error("Error closing "+projectConfig.Out, err)
cmdLogger.Error("Error closing contracts.json", err)
}
}(file)

// Write to the file
_, err = file.WriteString(invariants)
if err != nil {
return err
}
// Start the API to serve
api.Start(projectConfig, slitherOutput)

return nil
}
31 changes: 0 additions & 31 deletions cmd/generate_flags.go
Original file line number Diff line number Diff line change
Expand Up @@ -24,13 +24,6 @@ func addGenerateFlags() error {
// Flags
generateCmd.Flags().String("out", "",
fmt.Sprintf("path to output directory (unless a config file is provided, default is %q)", defaultConfig.Out))
generateCmd.Flags().String("target-contracts-dir", "",
fmt.Sprintf("directory path for target contracts (unless a config file is provided, default is %q)", defaultConfig.TargetContracts.Dir))
generateCmd.Flags().String("fuzz-tests-dir", "", fmt.Sprintf("directory path for fuzz tests (unless a config file is provided, default is %q)", defaultConfig.FuzzTests.Dir))
generateCmd.Flags().String("unit-tests-dir", "",
fmt.Sprintf("directory path for unit tests (unless a config file is provided, default is %q)", defaultConfig.UnitTests.Dir))
generateCmd.Flags().String("coverage-report-file", "",
fmt.Sprintf("directory path for coverage report (unless a config file is provided, default is %q)", defaultConfig.CoverageReportFile))

return nil
}
Expand All @@ -55,29 +48,5 @@ func updateProjectConfigWithGenerateFlags(cmd *cobra.Command, projectConfig *con
}
}

// Update fuzz tests directory
if cmd.Flags().Changed("fuzz-tests-dir") {
projectConfig.FuzzTests.Dir, err = cmd.Flags().GetString("fuzz-tests-dir")
if err != nil {
return err
}
}

// Update unit tests directory
if cmd.Flags().Changed("unit-tests-dir") {
projectConfig.UnitTests.Dir, err = cmd.Flags().GetString("unit-tests-dir")
if err != nil {
return err
}
}

// Update coverage report file
if cmd.Flags().Changed("coverage-report-file") {
projectConfig.CoverageReportFile, err = cmd.Flags().GetString("coverage-report-file")
if err != nil {
return err
}
}

return nil
}
23 changes: 0 additions & 23 deletions cmd/init_flags.go
Original file line number Diff line number Diff line change
Expand Up @@ -28,28 +28,5 @@ func updateProjectConfigWithInitFlags(cmd *cobra.Command, projectConfig *config.
}
}

// Update fuzz tests directory
if cmd.Flags().Changed("fuzz-tests-dir") {
projectConfig.FuzzTests.Dir, err = cmd.Flags().GetString("fuzz-tests-dir")
if err != nil {
return err
}
}

// Update unit tests directory
if cmd.Flags().Changed("unit-tests-dir") {
projectConfig.UnitTests.Dir, err = cmd.Flags().GetString("unit-tests-dir")
if err != nil {
return err
}
}

// Update coverage report file
if cmd.Flags().Changed("coverage-report-file") {
projectConfig.CoverageReportFile, err = cmd.Flags().GetString("coverage-report-file")
if err != nil {
return err
}
}
return nil
}
13 changes: 2 additions & 11 deletions config/config.go
Original file line number Diff line number Diff line change
Expand Up @@ -21,17 +21,8 @@ type ProjectConfig struct {
// TargetContracts describes the directory that holds the contracts to be fuzzed.
TargetContracts DirectoryConfig `json:"targetContracts"`

// FuzzTests describes the directory that holds the fuzz tests.
FuzzTests DirectoryConfig `json:"fuzzTests"`

// UnitTests describes the directory that holds the unit tests.
UnitTests DirectoryConfig `json:"unitTests"`

// CoverageReportFile describes the path to the coverage report file
CoverageReportFile string `json:"coverageReportFile"`

// SupportingFiles describes the paths to the files that provide additional information about the codebase
SupportingFiles []string `json:"supportingFiles"`
// Port describes the port that the API will be running on
Port int `json:"port"`

// NumInvariants describes the number of invariants to generate at a time
NumInvariants int `json:"numInvariants"`
Expand Down
12 changes: 1 addition & 11 deletions config/config_defaults.go
Original file line number Diff line number Diff line change
Expand Up @@ -10,17 +10,7 @@ func GetDefaultProjectConfig() (*ProjectConfig, error) {
Dir: "",
ExcludePaths: []string{},
},
FuzzTests: DirectoryConfig{
Dir: "",
ExcludePaths: []string{},
},
UnitTests: DirectoryConfig{
Dir: "",
ExcludePaths: []string{},
},
SupportingFiles: []string{},
NumInvariants: 10,
CoverageReportFile: "corpus/coverage_report.html",
NumInvariants: 10,
}

// Return the project configuration
Expand Down
1 change: 1 addition & 0 deletions go.mod
Original file line number Diff line number Diff line change
Expand Up @@ -12,6 +12,7 @@ require (
require (
github.com/dlclark/regexp2 v1.10.0 // indirect
github.com/google/uuid v1.3.0 // indirect
github.com/gorilla/mux v1.8.1 // indirect
github.com/inconshreveable/mousetrap v1.1.0 // indirect
github.com/mattn/go-colorable v0.1.13 // indirect
github.com/mattn/go-isatty v0.0.20 // indirect
Expand Down
2 changes: 2 additions & 0 deletions go.sum
Original file line number Diff line number Diff line change
Expand Up @@ -5,6 +5,8 @@ github.com/dlclark/regexp2 v1.10.0/go.mod h1:DHkYz0B9wPfa6wondMfaivmHpzrQ3v9q8cn
github.com/godbus/dbus/v5 v5.0.4/go.mod h1:xhWf0FNVPg57R7Z0UbKHbJfkEywrmjJnf7w5xrFpKfA=
github.com/google/uuid v1.3.0 h1:t6JiXgmwXMjEs8VusXIJk2BXHsn+wx8BZdTaoZ5fu7I=
github.com/google/uuid v1.3.0/go.mod h1:TIyPZe4MgqvfeYDBFedMoGGpEw/LqOeaOT+nhxU+yHo=
github.com/gorilla/mux v1.8.1 h1:TuBL49tXwgrFYWhqrNgrUNEY92u81SPhu7sTdzQEiWY=
github.com/gorilla/mux v1.8.1/go.mod h1:AKf9I4AEqPTmMytcMc0KkNouC66V3BtZ4qD5fmWSiMQ=
github.com/inconshreveable/mousetrap v1.1.0 h1:wN+x4NVGpMsO7ErUn/mUI3vEoE6Jt13X2s0bqwp9tc8=
github.com/inconshreveable/mousetrap v1.1.0/go.mod h1:vpF70FUmC8bwa3OWnCshd2FqLfsEA9PFc4w1p2J65bw=
github.com/mattn/go-colorable v0.1.13 h1:fFA4WZxdEF4tXPZVKMLwD8oUnCTTo08duU7wxecdEvA=
Expand Down
Loading

0 comments on commit 6edd87c

Please sign in to comment.