-
Notifications
You must be signed in to change notification settings - Fork 838
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
Merge pull request #1762 from projectdiscovery/http_based_speed_control
add http based speed control
- Loading branch information
Showing
8 changed files
with
217 additions
and
16 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,99 @@ | ||
| package main | ||
|
|
||
| import ( | ||
| "bytes" | ||
| "encoding/json" | ||
| "fmt" | ||
| "log" | ||
| "net/http" | ||
| "time" | ||
|
|
||
| "github.com/projectdiscovery/goflags" | ||
| "github.com/projectdiscovery/gologger" | ||
| "github.com/projectdiscovery/gologger/levels" | ||
| "github.com/projectdiscovery/httpx/runner" | ||
| ) | ||
|
|
||
| func main() { | ||
| gologger.DefaultLogger.SetMaxLevel(levels.LevelVerbose) // increase the verbosity (optional) | ||
|
|
||
| // generate urls | ||
| var urls []string | ||
| for i := 0; i < 100; i++ { | ||
| urls = append(urls, fmt.Sprintf("https://scanme.sh/a=%d", i)) | ||
| } | ||
|
|
||
| apiEndpoint := "127.0.0.1:31234" | ||
|
|
||
| options := runner.Options{ | ||
| Methods: "GET", | ||
| InputTargetHost: goflags.StringSlice(urls), | ||
| Threads: 1, | ||
| HttpApiEndpoint: apiEndpoint, | ||
| OnResult: func(r runner.Result) { | ||
| // handle error | ||
| if r.Err != nil { | ||
| fmt.Printf("[Err] %s: %s\n", r.Input, r.Err) | ||
| return | ||
| } | ||
| fmt.Printf("%s %s %d\n", r.Input, r.Host, r.StatusCode) | ||
| }, | ||
| } | ||
|
|
||
| // after 3 seconds increase the speed to 50 | ||
| time.AfterFunc(3*time.Second, func() { | ||
| client := &http.Client{} | ||
|
|
||
| concurrencySettings := runner.Concurrency{Threads: 50} | ||
| requestBody, err := json.Marshal(concurrencySettings) | ||
| if err != nil { | ||
| log.Fatalf("Error creating request body: %v", err) | ||
| } | ||
|
|
||
| req, err := http.NewRequest("PUT", fmt.Sprintf("http://%s/api/concurrency", apiEndpoint), bytes.NewBuffer(requestBody)) | ||
| if err != nil { | ||
| log.Fatalf("Error creating PUT request: %v", err) | ||
| } | ||
| req.Header.Set("Content-Type", "application/json") | ||
|
|
||
| resp, err := client.Do(req) | ||
| if err != nil { | ||
| log.Fatalf("Error sending PUT request: %v", err) | ||
| } | ||
| defer resp.Body.Close() | ||
|
|
||
| if resp.StatusCode != http.StatusOK { | ||
| log.Printf("Failed to update threads, status code: %d", resp.StatusCode) | ||
| } else { | ||
| log.Println("Threads updated to 50 successfully") | ||
| } | ||
| }) | ||
|
|
||
| if err := options.ValidateOptions(); err != nil { | ||
| log.Fatal(err) | ||
| } | ||
|
|
||
| httpxRunner, err := runner.New(&options) | ||
| if err != nil { | ||
| log.Fatal(err) | ||
| } | ||
| defer httpxRunner.Close() | ||
|
|
||
| httpxRunner.RunEnumeration() | ||
|
|
||
| // check the threads | ||
| req, err := http.Get(fmt.Sprintf("http://%s/api/concurrency", apiEndpoint)) | ||
| if err != nil { | ||
| log.Fatalf("Error creating GET request: %v", err) | ||
| } | ||
| var concurrencySettings runner.Concurrency | ||
| if err := json.NewDecoder(req.Body).Decode(&concurrencySettings); err != nil { | ||
| log.Fatalf("Error decoding response body: %v", err) | ||
| } | ||
|
|
||
| if concurrencySettings.Threads == 50 { | ||
| log.Println("Threads are set to 50") | ||
| } else { | ||
| log.Fatalf("Fatal error: Threads are not set to 50, current value: %d", concurrencySettings.Threads) | ||
| } | ||
| } |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,73 @@ | ||
| // TODO: move this to internal package | ||
| package runner | ||
|
|
||
| import ( | ||
| "encoding/json" | ||
| "net/http" | ||
| ) | ||
|
|
||
| type Concurrency struct { | ||
| Threads int `json:"threads"` | ||
| } | ||
|
|
||
| // Server represents the HTTP server that handles the concurrency settings endpoints. | ||
| type Server struct { | ||
| addr string | ||
| config *Options | ||
| } | ||
|
|
||
| // New creates a new instance of Server. | ||
| func NewServer(addr string, config *Options) *Server { | ||
| return &Server{ | ||
| addr: addr, | ||
| config: config, | ||
| } | ||
| } | ||
|
|
||
| // Start initializes the server and its routes, then starts listening on the specified address. | ||
| func (s *Server) Start() error { | ||
| http.HandleFunc("/api/concurrency", s.handleConcurrency) | ||
| if err := http.ListenAndServe(s.addr, nil); err != nil { | ||
| return err | ||
| } | ||
| return nil | ||
| } | ||
|
|
||
| // handleConcurrency routes the request based on its method to the appropriate handler. | ||
| func (s *Server) handleConcurrency(w http.ResponseWriter, r *http.Request) { | ||
| switch r.Method { | ||
| case http.MethodGet: | ||
| s.getSettings(w, r) | ||
| case http.MethodPut: | ||
| s.updateSettings(w, r) | ||
| default: | ||
| http.Error(w, "Unsupported HTTP method", http.StatusMethodNotAllowed) | ||
| } | ||
| } | ||
|
|
||
| // GetSettings handles GET requests and returns the current concurrency settings | ||
| func (s *Server) getSettings(w http.ResponseWriter, _ *http.Request) { | ||
| concurrencySettings := Concurrency{ | ||
| Threads: s.config.Threads, | ||
| } | ||
| w.Header().Set("Content-Type", "application/json") | ||
| if err := json.NewEncoder(w).Encode(concurrencySettings); err != nil { | ||
| http.Error(w, err.Error(), http.StatusInternalServerError) | ||
| return | ||
| } | ||
| } | ||
|
|
||
| // UpdateSettings handles PUT requests to update the concurrency settings | ||
| func (s *Server) updateSettings(w http.ResponseWriter, r *http.Request) { | ||
| var newSettings Concurrency | ||
| if err := json.NewDecoder(r.Body).Decode(&newSettings); err != nil { | ||
| http.Error(w, err.Error(), http.StatusBadRequest) | ||
| return | ||
| } | ||
|
|
||
| if newSettings.Threads > 0 { | ||
| s.config.Threads = newSettings.Threads | ||
| } | ||
|
|
||
| w.WriteHeader(http.StatusOK) | ||
| } |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters