Skip to content

Commit 7d050cf

Browse files
VitaliiMaltsevVitaliiMaltsev
authored
Postgres Source Strict Encrypt: Add tests with SSL and SSH (#19388)
* Postgres Source Strict Encrypt: Add tests with SSL and SSH * updated No tunnel test * fixed spec test * fixed spec test * refactoring
1 parent a1db24c commit 7d050cf

File tree

2 files changed

+100
-0
lines changed

2 files changed

+100
-0
lines changed

‎airbyte-integrations/connectors/source-postgres/src/main/java/io/airbyte/integrations/source/postgres/PostgresSourceStrictEncrypt.java

+33
Original file line numberDiff line numberDiff line change
@@ -4,24 +4,36 @@
44

55
package io.airbyte.integrations.source.postgres;
66

7+
import com.fasterxml.jackson.databind.JsonNode;
78
import com.fasterxml.jackson.databind.node.ArrayNode;
89
import com.fasterxml.jackson.databind.node.ObjectNode;
910
import io.airbyte.commons.json.Jsons;
1011
import io.airbyte.db.jdbc.JdbcUtils;
1112
import io.airbyte.integrations.base.IntegrationRunner;
1213
import io.airbyte.integrations.base.Source;
1314
import io.airbyte.integrations.base.spec_modification.SpecModifyingSource;
15+
import io.airbyte.protocol.models.AirbyteConnectionStatus;
1416
import io.airbyte.protocol.models.ConnectorSpecification;
17+
import java.util.Set;
1518
import org.slf4j.Logger;
1619
import org.slf4j.LoggerFactory;
1720

21+
import static io.airbyte.protocol.models.AirbyteConnectionStatus.Status;
22+
1823
/**
1924
* This class is copied from source-postgres-strict-encrypt. The original file can be deleted
2025
* completely once the migration of multi-variant connector is done.
2126
*/
2227
public class PostgresSourceStrictEncrypt extends SpecModifyingSource implements Source {
2328

2429
private static final Logger LOGGER = LoggerFactory.getLogger(PostgresSourceStrictEncrypt.class);
30+
public static final String TUNNEL_METHOD = "tunnel_method";
31+
public static final String NO_TUNNEL = "NO_TUNNEL";
32+
public static final String SSL_MODE = "ssl_mode";
33+
public static final String MODE = "mode";
34+
public static final String SSL_MODE_ALLOW = "allow";
35+
public static final String SSL_MODE_PREFER = "prefer";
36+
public static final String SSL_MODE_DISABLE = "disable";
2537

2638
public PostgresSourceStrictEncrypt() {
2739
super(PostgresSource.sshWrappedSource());
@@ -39,6 +51,27 @@ public ConnectorSpecification modifySpec(final ConnectorSpecification originalSp
3951
return spec;
4052
}
4153

54+
@Override
55+
public AirbyteConnectionStatus check(final JsonNode config) throws Exception {
56+
// #15808 Disallow connecting to db with disable, prefer or allow SSL mode when connecting directly
57+
// and not over SSH tunnel
58+
if (config.has(TUNNEL_METHOD)
59+
&& config.get(TUNNEL_METHOD).has(TUNNEL_METHOD)
60+
&& config.get(TUNNEL_METHOD).get(TUNNEL_METHOD).asText().equals(NO_TUNNEL)) {
61+
// If no SSH tunnel
62+
if (config.has(SSL_MODE) && config.get(SSL_MODE).has(MODE)) {
63+
if (Set.of(SSL_MODE_DISABLE, SSL_MODE_ALLOW, SSL_MODE_PREFER).contains(config.get(SSL_MODE).get(MODE).asText())) {
64+
// Fail in case SSL mode is disable, allow or prefer
65+
return new AirbyteConnectionStatus()
66+
.withStatus(Status.FAILED)
67+
.withMessage(
68+
"Unsecured connection not allowed. If no SSH Tunnel set up, please use one of the following SSL modes: require, verify-ca, verify-full");
69+
}
70+
}
71+
}
72+
return super.check(config);
73+
}
74+
4275
public static void main(final String[] args) throws Exception {
4376
final Source source = new PostgresSourceStrictEncrypt();
4477
LOGGER.info("starting source: {}", PostgresSourceStrictEncrypt.class);

‎airbyte-integrations/connectors/source-postgres/src/test/java/io/airbyte/integrations/source/postgres/PostgresSourceStrictEncryptTest.java

+67
Original file line numberDiff line numberDiff line change
@@ -5,9 +5,12 @@
55
package io.airbyte.integrations.source.postgres;
66

77
import static org.junit.jupiter.api.Assertions.assertEquals;
8+
import static org.junit.jupiter.api.Assertions.assertTrue;
89

910
import com.fasterxml.jackson.databind.JsonNode;
11+
import com.fasterxml.jackson.databind.node.ObjectNode;
1012
import com.google.common.collect.ImmutableMap;
13+
import io.airbyte.commons.json.Jsons;
1114
import io.airbyte.db.jdbc.JdbcUtils;
1215
import io.airbyte.integrations.base.ssh.SshBastionContainer;
1316
import io.airbyte.integrations.base.ssh.SshTunnel;
@@ -18,6 +21,7 @@
1821
import org.junit.jupiter.api.Test;
1922
import org.testcontainers.containers.Network;
2023
import org.testcontainers.containers.PostgreSQLContainer;
24+
import org.testcontainers.utility.DockerImageName;
2125

2226
public class PostgresSourceStrictEncryptTest {
2327

@@ -76,4 +80,67 @@ private ImmutableMap.Builder<Object, Object> getDatabaseConfigBuilderWithSSLMode
7680
.put(JdbcUtils.SSL_MODE_KEY, Map.of(JdbcUtils.MODE_KEY, sslMode));
7781
}
7882

83+
private JsonNode getMockedSSLConfig(String sslMode) {
84+
return Jsons.jsonNode(ImmutableMap.builder()
85+
.put(JdbcUtils.HOST_KEY, "test_host")
86+
.put(JdbcUtils.PORT_KEY, 777)
87+
.put(JdbcUtils.DATABASE_KEY, "test_db")
88+
.put(JdbcUtils.USERNAME_KEY, "test_user")
89+
.put(JdbcUtils.PASSWORD_KEY, "test_password")
90+
.put(JdbcUtils.SSL_KEY, true)
91+
.put(JdbcUtils.SSL_MODE_KEY, Map.of(JdbcUtils.MODE_KEY, sslMode))
92+
.build());
93+
}
94+
95+
@Test
96+
void testSslModesUnsecuredNoTunnel() throws Exception {
97+
for (String sslMode : List.of("disable", "allow", "prefer")) {
98+
final JsonNode config = getMockedSSLConfig(sslMode);
99+
((ObjectNode) config).putIfAbsent("tunnel_method", Jsons.jsonNode(ImmutableMap.builder()
100+
.put("tunnel_method", "NO_TUNNEL")
101+
.build()));
102+
103+
final AirbyteConnectionStatus actual = new PostgresSourceStrictEncrypt().check(config);
104+
assertEquals(AirbyteConnectionStatus.Status.FAILED, actual.getStatus());
105+
assertTrue(actual.getMessage().contains("Unsecured connection not allowed"));
106+
}
107+
}
108+
109+
@Test
110+
void testSslModeRequiredNoTunnel() throws Exception {
111+
112+
try (PostgreSQLContainer<?> db =
113+
new PostgreSQLContainer<>(DockerImageName.parse("marcosmarxm/postgres-ssl:dev").asCompatibleSubstituteFor("postgres"))
114+
.withCommand("postgres -c ssl=on -c ssl_cert_file=/var/lib/postgresql/server.crt -c ssl_key_file=/var/lib/postgresql/server.key")) {
115+
db.start();
116+
117+
final ImmutableMap<Object, Object> configBuilderWithSslModeRequire = getDatabaseConfigBuilderWithSSLMode(db, "require").build();
118+
final JsonNode config = Jsons.jsonNode(configBuilderWithSslModeRequire);
119+
((ObjectNode) config).putIfAbsent("tunnel_method", Jsons.jsonNode(ImmutableMap.builder()
120+
.put("tunnel_method", "NO_TUNNEL")
121+
.build()));
122+
final AirbyteConnectionStatus connectionStatusForPreferredMode = new PostgresSourceStrictEncrypt().check(config);
123+
assertEquals(AirbyteConnectionStatus.Status.SUCCEEDED, connectionStatusForPreferredMode.getStatus());
124+
}
125+
}
126+
127+
@Test
128+
void testStrictSSLSecuredWithTunnel() throws Exception {
129+
try (PostgreSQLContainer<?> db =
130+
new PostgreSQLContainer<>(DockerImageName.parse("marcosmarxm/postgres-ssl:dev").asCompatibleSubstituteFor("postgres"))
131+
.withCommand("postgres -c ssl=on -c ssl_cert_file=/var/lib/postgresql/server.crt -c ssl_key_file=/var/lib/postgresql/server.key")
132+
.withNetwork(network)) {
133+
134+
bastion.initAndStartBastion(network);
135+
db.start();
136+
137+
final ImmutableMap.Builder<Object, Object> builderWithSSLModePrefer = getDatabaseConfigBuilderWithSSLMode(db, "require");
138+
final JsonNode configWithSslAndSsh = bastion.getTunnelConfig(SshTunnel.TunnelMethod.SSH_PASSWORD_AUTH, builderWithSSLModePrefer);
139+
final AirbyteConnectionStatus connectionStatusForPreferredMode = new PostgresSourceStrictEncrypt().check(configWithSslAndSsh);
140+
assertEquals(AirbyteConnectionStatus.Status.SUCCEEDED, connectionStatusForPreferredMode.getStatus());
141+
} finally {
142+
bastion.stopAndClose();
143+
}
144+
}
145+
79146
}

0 commit comments

Comments
 (0)