Skip to content
This repository was archived by the owner on Dec 3, 2024. It is now read-only.

Commit 5662d55

Browse files
cimninecimnine
committed
Implement rate limiter
… and fix tests
1 parent 6e9a468 commit 5662d55

File tree

7 files changed

+142
-40
lines changed

7 files changed

+142
-40
lines changed

‎.editorconfig‎

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -8,6 +8,13 @@ indent_style = space
88
insert_final_newline = true
99
trim_trailing_whitespace = true
1010

11+
[{Makefile,go.mod,go.sum,*.go,.gitmodules}]
12+
indent_style = tab
13+
indent_size = 4
14+
15+
[*.md]
16+
trim_trailing_whitespace = false
17+
1118
[GNUmakefile]
1219
indent_style = tab
1320
tab_width = 4

‎go.mod‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -3,10 +3,12 @@ module gitlab.com/nxt/public/terraform-provider-publicip
33
go 1.17
44

55
require (
6+
github.com/hashicorp/go-uuid v1.0.2
67
github.com/hashicorp/terraform-plugin-docs v0.5.1
78
github.com/hashicorp/terraform-plugin-framework v0.5.0
89
github.com/hashicorp/terraform-plugin-go v0.6.0
910
github.com/hashicorp/terraform-plugin-sdk/v2 v2.10.1
11+
golang.org/x/time v0.0.0-20191024005414-555d28b269f0
1012
)
1113

1214
require (
@@ -38,7 +40,6 @@ require (
3840
github.com/hashicorp/go-multierror v1.1.1 // indirect
3941
github.com/hashicorp/go-plugin v1.4.3 // indirect
4042
github.com/hashicorp/go-safetemp v1.0.0 // indirect
41-
github.com/hashicorp/go-uuid v1.0.2 // indirect
4243
github.com/hashicorp/go-version v1.3.0 // indirect
4344
github.com/hashicorp/hc-install v0.3.1 // indirect
4445
github.com/hashicorp/hcl/v2 v2.3.0 // indirect

‎go.sum‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -500,6 +500,7 @@ golang.org/x/text v0.3.5 h1:i6eZZ+zk0SOf0xgBpEpPD18qWcJda6q1sxt3S0kzyUQ=
500500
golang.org/x/text v0.3.5/go.mod h1:5Zoc/QRtKVWzQhOtBMvqHzDpF6irO9z98xDceosuGiQ=
501501
golang.org/x/time v0.0.0-20181108054448-85acf8d2951c/go.mod h1:tRJNPiyCQ0inRvYxbN9jk5I+vvW/OXSQhTDSoE431IQ=
502502
golang.org/x/time v0.0.0-20190308202827-9d24e82272b4/go.mod h1:tRJNPiyCQ0inRvYxbN9jk5I+vvW/OXSQhTDSoE431IQ=
503+
golang.org/x/time v0.0.0-20191024005414-555d28b269f0 h1:/5xXl8Y5W96D+TtHSlonuFqGHIWVuyCkGJLwGh9JJFs=
503504
golang.org/x/time v0.0.0-20191024005414-555d28b269f0/go.mod h1:tRJNPiyCQ0inRvYxbN9jk5I+vvW/OXSQhTDSoE431IQ=
504505
golang.org/x/tools v0.0.0-20180917221912-90fa682c2a6e/go.mod h1:n7NCudcB/nEzxVGmLbDWY5pfWTLqBcC2KZ6jyYvM4mQ=
505506
golang.org/x/tools v0.0.0-20190114222345-bf090417da8b/go.mod h1:n7NCudcB/nEzxVGmLbDWY5pfWTLqBcC2KZ6jyYvM4mQ=

‎internal/provider/ip_address_data_source.go‎

Lines changed: 33 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -10,23 +10,29 @@ import (
1010
"path"
1111
"strings"
1212

13+
"github.com/hashicorp/go-uuid"
1314
"github.com/hashicorp/terraform-plugin-framework/diag"
1415
"github.com/hashicorp/terraform-plugin-framework/tfsdk"
1516
"github.com/hashicorp/terraform-plugin-framework/types"
1617
)
1718

1819
type ipDataSourceType struct{}
1920

20-
func (t ipDataSourceType) GetSchema(ctx context.Context) (tfsdk.Schema, diag.Diagnostics) {
21+
func (t ipDataSourceType) GetSchema(_ context.Context) (tfsdk.Schema, diag.Diagnostics) {
2122
return tfsdk.Schema{
2223
// This description is used by the documentation generator and the language server.
2324
MarkdownDescription: "The current (public) IP as reported by the IP information provider.",
2425

2526
Attributes: map[string]tfsdk.Attribute{
26-
"ip_version": {
27+
"id": {
28+
MarkdownDescription: "An ID, which is only used internally. *Do not use this field in your terraform definitions.*",
29+
Computed: true,
2730
Type: types.StringType,
31+
},
32+
"ip_version": {
2833
MarkdownDescription: "Whether to use IPv4 or IPv6 only. Valid values: 'V4', 'V6'",
2934
Optional: true,
35+
Type: types.StringType,
3036
Validators: []tfsdk.AttributeValidator{ipVersionValidator{}},
3137
},
3238
"ip": {
@@ -48,7 +54,7 @@ func (t ipDataSourceType) GetSchema(ctx context.Context) (tfsdk.Schema, diag.Dia
4854
}, nil
4955
}
5056

51-
func (t ipDataSourceType) NewDataSource(ctx context.Context, in tfsdk.Provider) (tfsdk.DataSource, diag.Diagnostics) {
57+
func (t ipDataSourceType) NewDataSource(_ context.Context, in tfsdk.Provider) (tfsdk.DataSource, diag.Diagnostics) {
5258
provider, diags := convertProviderType(in)
5359

5460
return ipDataSource{
@@ -57,6 +63,7 @@ func (t ipDataSourceType) NewDataSource(ctx context.Context, in tfsdk.Provider)
5763
}
5864

5965
type ipDataSourceData struct {
66+
ID types.String `tfsdk:"id"`
6067
IPVersion types.String `tfsdk:"ip_version"`
6168
IP types.String `tfsdk:"ip"`
6269
ASNID types.String `tfsdk:"asn_id"`
@@ -121,9 +128,21 @@ func (d ipDataSource) Read(ctx context.Context, req tfsdk.ReadDataSourceRequest,
121128

122129
log.Printf("got to send request ✅: %s", userAgent)
123130

131+
if !d.provider.rateLimiter.Allow() {
132+
log.Printf("the rate limit may be triggered ⏳")
133+
}
134+
135+
timeoutCtx, cancelFunc := context.WithTimeout(ctx, d.provider.timeout)
136+
defer cancelFunc()
137+
err = d.provider.rateLimiter.Wait(timeoutCtx)
138+
if err != nil {
139+
log.Printf("Rate limiter error 🚨: %s", err)
140+
resp.Diagnostics.AddError("Error waiting for rate limit", fmt.Sprintf("There was an error while awaiting a slot from the rate limiter: %s", err))
141+
}
142+
124143
httpResp, err := client.Do(httpReq)
125144
if err != nil {
126-
log.Printf("HTTP Client Error 🚨: %s", err)
145+
log.Printf("HTTP client error 🚨: %s", err)
127146
resp.Diagnostics.AddError("Error fetching information from the IP information provider", fmt.Sprintf("There was an error when contacting '%s': %s", requestURLstr, err))
128147
return
129148
}
@@ -151,6 +170,16 @@ func (d ipDataSource) Read(ctx context.Context, req tfsdk.ReadDataSourceRequest,
151170

152171
log.Printf("got to apply ✅: %+v", respData)
153172

173+
if data.ID.Unknown || data.ID.Null {
174+
uuidStr, err := uuid.GenerateUUID()
175+
if err != nil {
176+
log.Printf("Error while generating a new UUID 🚨: %s", err)
177+
resp.Diagnostics.AddError("Internal error, try again.", fmt.Sprintf("There was an internal error in the provider when creating a new UUID: %s.", err))
178+
return
179+
}
180+
data.ID = types.String{Value: uuidStr}
181+
}
182+
154183
data.IP = types.String{Value: respData.IP}
155184
data.ASNID = types.String{Value: respData.ASN}
156185
data.ASNOrg = types.String{Value: respData.ASNOrg}

‎internal/provider/ip_address_data_source_test.go‎

Lines changed: 7 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -11,23 +11,22 @@ func TestAccExampleDataSource(t *testing.T) {
1111
PreCheck: func() { testAccPreCheck(t) },
1212
ProtoV6ProviderFactories: testAccProtoV6ProviderFactories,
1313
Steps: []resource.TestStep{
14-
// Read testing
1514
{
16-
Config: v4Config,
15+
Config: defaultConfig,
1716
Check: resource.ComposeAggregateTestCheckFunc(
18-
resource.TestCheckResourceAttr("data.publicip_address.default", "ip_version", "v4"),
17+
resource.TestCheckResourceAttrSet("data.publicip_address.default", "ip_version"),
1918
),
2019
},
2120
{
2221
Config: v6Config,
2322
Check: resource.ComposeAggregateTestCheckFunc(
24-
resource.TestCheckResourceAttr("data.publicip_address.default", "ip_version", "v6"),
23+
resource.TestCheckResourceAttr("data.publicip_address.v6", "ip_version", "v6"),
2524
),
2625
},
2726
{
28-
Config: defaultConfig,
27+
Config: v4Config,
2928
Check: resource.ComposeAggregateTestCheckFunc(
30-
resource.TestCheckResourceAttrSet("data.publicip_address.default", "ip_version"),
29+
resource.TestCheckResourceAttr("data.publicip_address.v4", "ip_version", "v4"),
3130
),
3231
},
3332
},
@@ -41,12 +40,12 @@ data "publicip_address" "default" {
4140

4241
const v4Config = `
4342
data "publicip_address" "v4" {
44-
ip_version = "v4"
43+
ip_version = "v4"
4544
}
4645
`
4746

4847
const v6Config = `
4948
data "publicip_address" "v6" {
50-
ip_version = "v6"
49+
ip_version = "v6"
5150
}
5251
`

‎internal/provider/ip_response.go‎

Lines changed: 18 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -1,20 +1,24 @@
11
package provider
22

3+
import (
4+
"encoding/json"
5+
)
6+
37
type IPResponse struct {
4-
IP string `json:"ip,omitempty"`
5-
IPDecimal int64 `json:"ip_decimal,omitempty"`
6-
Country string `json:"country,omitempty"`
7-
CountryISO string `json:"country_iso,omitempty"`
8-
CountryEU bool `json:"country_eu,omitempty"`
9-
RegionName string `json:"region_name,omitempty"`
10-
RegionCode string `json:"region_code,omitempty"`
11-
ZIPCode string `json:"zip_code,omitempty"`
12-
City string `json:"city,omitempty"`
13-
Latitude float32 `json:"latitude,omitempty"`
14-
Longitude float32 `json:"longitude,omitempty"`
15-
TimeZone string `json:"time_zone,omitempty"`
16-
ASN string `json:"asn,omitempty"`
17-
ASNOrg string `json:"asn_org,omitempty"`
8+
IP string `json:"ip,omitempty"`
9+
IPDecimal json.Number `json:"ip_decimal,omitempty"`
10+
Country string `json:"country,omitempty"`
11+
CountryISO string `json:"country_iso,omitempty"`
12+
CountryEU bool `json:"country_eu,omitempty"`
13+
RegionName string `json:"region_name,omitempty"`
14+
RegionCode string `json:"region_code,omitempty"`
15+
ZIPCode string `json:"zip_code,omitempty"`
16+
City string `json:"city,omitempty"`
17+
Latitude float32 `json:"latitude,omitempty"`
18+
Longitude float32 `json:"longitude,omitempty"`
19+
TimeZone string `json:"time_zone,omitempty"`
20+
ASN string `json:"asn,omitempty"`
21+
ASNOrg string `json:"asn_org,omitempty"`
1822
UserAgent struct {
1923
Product string `json:"product,omitempty"`
2024
Version string `json:"version,omitempty"`

‎internal/provider/provider.go‎

Lines changed: 74 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -3,19 +3,22 @@ package provider
33
import (
44
"context"
55
"fmt"
6+
"math"
67
"net/url"
78
"time"
89

910
"github.com/hashicorp/terraform-plugin-framework/diag"
1011
"github.com/hashicorp/terraform-plugin-framework/tfsdk"
1112
"github.com/hashicorp/terraform-plugin-framework/types"
13+
"golang.org/x/time/rate"
1214
)
1315

1416
// provider satisfies the tfsdk.Provider interface and usually is included
1517
// with all Resource and DataSource implementations.
1618
type provider struct {
17-
timeout time.Duration
18-
ipURL *url.URL
19+
timeout time.Duration
20+
ipURL *url.URL
21+
rateLimiter *rate.Limiter
1922

2023
// configured is set to true at the end of the Configure method.
2124
// This can be used in Resource and DataSource implementations to verify
@@ -33,12 +36,16 @@ type provider struct {
3336

3437
// providerData can be used to store data from the Terraform configuration.
3538
type providerData struct {
36-
ProviderURL types.String `tfsdk:"provider_url"`
37-
Timeout types.String `tfsdk:"timeout"`
39+
ProviderURL types.String `tfsdk:"provider_url"`
40+
Timeout types.String `tfsdk:"timeout"`
41+
RateLimitRate types.String `tfsdk:"rate_limit_rate"`
42+
RateLimitBurst types.Int64 `tfsdk:"rate_limit_burst"`
3843
}
3944

4045
const DefaultTimeout = "5s"
4146
const DefaultProviderURL = "https://ifconfig.co/"
47+
const DefaultRateLimitRate = "500ms"
48+
const DefaultRateLimitBurst = 1
4249

4350
func (p *provider) Configure(ctx context.Context, req tfsdk.ConfigureProviderRequest, resp *tfsdk.ConfigureProviderResponse) {
4451
var data providerData
@@ -56,33 +63,77 @@ func (p *provider) Configure(ctx context.Context, req tfsdk.ConfigureProviderReq
5663
providerURL = data.ProviderURL.Value
5764
}
5865

66+
if !p.configureProviderURL(providerURL, resp) ||
67+
!p.configureTimeout(data, resp) ||
68+
!p.configureRateLimiter(data, resp) {
69+
return
70+
}
71+
72+
p.configured = true
73+
}
74+
75+
func (p *provider) configureProviderURL(providerURL string, resp *tfsdk.ConfigureProviderResponse) bool {
5976
var err error
6077
p.ipURL, err = url.Parse(providerURL)
78+
6179
if err != nil {
62-
resp.Diagnostics.AddError("Unable to parse the provider_url", fmt.Sprintf("The provider_url '%s' can't be parsed: %s", providerURL, err))
63-
return
80+
resp.Diagnostics.AddError("Unable to parse the provider_url", fmt.Sprintf("The provider_url value '%s' can't be parsed: %s", providerURL, err))
81+
return false
6482
}
83+
return true
84+
}
6585

86+
func (p *provider) configureTimeout(data providerData, resp *tfsdk.ConfigureProviderResponse) bool {
6687
var timeout string
6788
if data.Timeout.Null {
6889
timeout = DefaultTimeout
6990
} else {
7091
timeout = data.Timeout.Value
7192
}
7293

94+
var err error
7395
p.timeout, err = time.ParseDuration(timeout)
7496
if err != nil {
75-
resp.Diagnostics.AddError("Unable to parse the timeout", fmt.Sprintf("The timeout '%s' can't be parsed: %s", timeout, err))
76-
return
97+
resp.Diagnostics.AddError("Unable to parse the timeout", fmt.Sprintf("The timeout value '%s' can't be parsed: %s", timeout, err))
98+
return false
7799
}
100+
return true
101+
}
78102

79-
p.configured = true
103+
func (p *provider) configureRateLimiter(data providerData, resp *tfsdk.ConfigureProviderResponse) bool {
104+
var rateLimitRate string
105+
if data.RateLimitRate.Null {
106+
rateLimitRate = DefaultRateLimitRate
107+
} else {
108+
rateLimitRate = data.RateLimitRate.Value
109+
}
110+
111+
rateLimitRateDuration, err := time.ParseDuration(rateLimitRate)
112+
if err != nil {
113+
resp.Diagnostics.AddError("Unable to parse the rate_limit_rate", fmt.Sprintf("The rate_limit_rate value '%s' can't be parsed: %s", rateLimitRate, err))
114+
return false
115+
}
116+
117+
var rateLimitBurst int
118+
if data.RateLimitBurst.Null {
119+
rateLimitBurst = DefaultRateLimitBurst
120+
} else if data.RateLimitBurst.Value > math.MaxInt {
121+
resp.Diagnostics.AddError("Unable to use the rate_limit_burst", fmt.Sprintf("The rate_limit_burst value '%d' is too big. Maximum allowed is %d", data.RateLimitBurst.Value, math.MaxInt))
122+
return false
123+
} else if data.RateLimitBurst.Value <= 0 {
124+
resp.Diagnostics.AddError("Unable to use the rate_limit_burst", fmt.Sprintf("The rate_limit_burst value '%d' must be bigger than 0", data.RateLimitBurst.Value))
125+
return false
126+
} else {
127+
rateLimitBurst = int(data.RateLimitBurst.Value)
128+
}
129+
130+
p.rateLimiter = rate.NewLimiter(rate.Every(rateLimitRateDuration), rateLimitBurst)
131+
132+
return true
80133
}
81134

82135
func (p *provider) GetResources(_ context.Context) (map[string]tfsdk.ResourceType, diag.Diagnostics) {
83-
return map[string]tfsdk.ResourceType{
84-
// "scaffolding_example": exampleResourceType{},
85-
}, nil
136+
return map[string]tfsdk.ResourceType{}, nil
86137
}
87138

88139
func (p *provider) GetDataSources(_ context.Context) (map[string]tfsdk.DataSourceType, diag.Diagnostics) {
@@ -95,10 +146,20 @@ func (p *provider) GetSchema(_ context.Context) (tfsdk.Schema, diag.Diagnostics)
95146
return tfsdk.Schema{
96147
Attributes: map[string]tfsdk.Attribute{
97148
"timeout": {
98-
MarkdownDescription: fmt.Sprintf("Timeout of the request to the IP information provider, defaults to `%s`.", DefaultTimeout),
149+
MarkdownDescription: fmt.Sprintf("Timeout of the request to the IP information provider. Defaults to `%s`.", DefaultTimeout),
99150
Optional: true,
100151
Type: types.StringType,
101152
},
153+
"rate_limit_rate": {
154+
MarkdownDescription: fmt.Sprintf("Limit the number of the request to the IP information provider. Defines the time until the limit is reset. Defaults to `%s`.", DefaultRateLimitRate),
155+
Optional: true,
156+
Type: types.StringType,
157+
},
158+
"rate_limit_burst": {
159+
MarkdownDescription: fmt.Sprintf("Limit the number of the request to the IP information provider. Defines the number of events per rate until the limit is reached. Defaults to `%d`.", DefaultRateLimitBurst),
160+
Optional: true,
161+
Type: types.Int64Type,
162+
},
102163
"provider_url": {
103164
MarkdownDescription: fmt.Sprintf("URL to a ifconfig.co-compatible IP information provider, defaults to `%s`.", DefaultProviderURL),
104165
Optional: true,

0 commit comments

Comments
 (0)